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": [],
"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();"
]
},

View File

@ -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,