fix device loading

rasbt · Jun 20, 2024 · 977f9c6 · 977f9c6
1 parent 8ccca3a
commit 977f9c6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/ch05/01_main-chapter-code/ch05.ipynb b/ch05/01_main-chapter-code/ch05.ipynb
@@ -1985,7 +1985,8 @@
    "outputs": [],
    "source": [
     "model = GPTModel(GPT_CONFIG_124M)\n",
-    "model.load_state_dict(torch.load(\"model.pth\"))\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "model.load_state_dict(torch.load(\"model.pth\", map_location=device))\n",
     "model.eval();"
    ]
   },

diff --git a/ch05/01_main-chapter-code/exercise-solutions.ipynb b/ch05/01_main-chapter-code/exercise-solutions.ipynb
@@ -427,6 +427,7 @@
     "checkpoint = torch.load(\"model_and_optimizer.pth\")\n",
     "model = GPTModel(GPT_CONFIG_124M)\n",
     "model.load_state_dict(checkpoint[\"model_state_dict\"])\n",
+    "model.to(device)\n",
     "\n",
     "optimizer = torch.optim.AdamW(model.parameters(), lr=0.0004, weight_decay=0.1)\n",
     "optimizer.load_state_dict(checkpoint[\"optimizer_state_dict\"])\n",
@@ -958,7 +959,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.4"
+   "version": "3.10.6"
   }
  },
  "nbformat": 4,