diff --git a/unit2/unit2.ipynb b/unit2/unit2.ipynb
index eefa1da..766b190 100644
--- a/unit2/unit2.ipynb
+++ b/unit2/unit2.ipynb
@@ -892,9 +892,9 @@
         "      total_rewards_ep += reward\n",
         "        \n",
         "      if done:\n",
-        "        episode_rewards.append(total_rewards_ep)\n",
         "        break\n",
         "      state = new_state\n",
+        "    episode_rewards.append(total_rewards_ep)\n",
         "  mean_reward = np.mean(episode_rewards)\n",
         "  std_reward = np.std(episode_rewards)\n",
         "\n",