diff --git a/week04_approx_rl/homework_pytorch_main.ipynb b/week04_approx_rl/homework_pytorch_main.ipynb
index d2628ec9..5dd1d253 100644
--- a/week04_approx_rl/homework_pytorch_main.ipynb
+++ b/week04_approx_rl/homework_pytorch_main.ipynb
@@ -475,7 +475,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def merge_frame_stack_to_plot(frame_stack_obs: np.ndarray | gym.wrappers.frame_stack.LazyFrames):\n",
+    "def merge_frame_stack_to_plot(frame_stack_obs: np.ndarray | gym.wrappers.stateful_observation.FrameStackObservation):\n",
     "    \"\"\"\n",
     "    A helper function to plot a frame stack as a single human-interpretable image.\n",
     "\n",