fix device loading

This commit is contained in:
rasbt 2024-06-20 08:07:00 -05:00
parent d440eb17bc
commit cb194fa8fa
2 changed files with 4 additions and 2 deletions

View File

@ -1985,7 +1985,8 @@
"outputs": [], "outputs": [],
"source": [ "source": [
"model = GPTModel(GPT_CONFIG_124M)\n", "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();" "model.eval();"
] ]
}, },

View File

@ -427,6 +427,7 @@
"checkpoint = torch.load(\"model_and_optimizer.pth\")\n", "checkpoint = torch.load(\"model_and_optimizer.pth\")\n",
"model = GPTModel(GPT_CONFIG_124M)\n", "model = GPTModel(GPT_CONFIG_124M)\n",
"model.load_state_dict(checkpoint[\"model_state_dict\"])\n", "model.load_state_dict(checkpoint[\"model_state_dict\"])\n",
"model.to(device)\n",
"\n", "\n",
"optimizer = torch.optim.AdamW(model.parameters(), lr=0.0004, weight_decay=0.1)\n", "optimizer = torch.optim.AdamW(model.parameters(), lr=0.0004, weight_decay=0.1)\n",
"optimizer.load_state_dict(checkpoint[\"optimizer_state_dict\"])\n", "optimizer.load_state_dict(checkpoint[\"optimizer_state_dict\"])\n",
@ -958,7 +959,7 @@
"name": "python", "name": "python",
"nbconvert_exporter": "python", "nbconvert_exporter": "python",
"pygments_lexer": "ipython3", "pygments_lexer": "ipython3",
"version": "3.11.4" "version": "3.10.6"
} }
}, },
"nbformat": 4, "nbformat": 4,