diff --git a/unit1/unit1.ipynb b/unit1/unit1.ipynb
index 105c902..709720a 100644
--- a/unit1/unit1.ipynb
+++ b/unit1/unit1.ipynb
@@ -695,7 +695,10 @@
       },
       "outputs": [],
       "source": [
-        "# Train it for 500,000 timesteps"
+        "# TODO: Train it for 500,000 timesteps\n",
+        "\n",
+        "# TODO: Specify file name for model and save the model to file\n",
+        "model_name = \"\"\n"
       ]
     },
     {
@@ -717,7 +720,10 @@
       "source": [
         "# SOLUTION\n",
         "# Train it for 500,000 timesteps\n",
-        "model.learn(total_timesteps=500000)"
+        "model.learn(total_timesteps=500000)\n",
+        "# Save the model\n",
+        "model_name = \"ppo-LunarLander-v2\"\n",
+        "model.save(model_name)"
       ]
     },
     {
@@ -980,9 +986,6 @@
         "# Define the name of the environment\n",
         "env_id = \"LunarLander-v2\"\n",
         "\n",
-        "# Define the name of the trained model that we defined in model_save\n",
-        "model_name = \"ppo-LunarLander-v2\"\n",
-        "\n",
         "# TODO: Define the model architecture we used\n",
         "model_architecture = \"PPO\"\n",
         "\n",