diff --git a/02_embedding_model.ipynb b/02_embedding_model.ipynb
index 01a3eff..c51fc20 100644
--- a/02_embedding_model.ipynb
+++ b/02_embedding_model.ipynb
@@ -1930,7 +1930,7 @@
    "source": [
     "import torch\n",
     "device = torch.device(\"mps\" if torch.cuda.is_available() else \"cpu\")\n",
-    "model.to(device)\n",
+    "improved_model.to(device)\n",
     "minicorpus_docs = minicorpus_docs.to(device)"
    ]
   },