Upload main.ipynb with huggingface_hub

main.ipynb

@@ -887,9 +887,8 @@
     "print(f\"Checkpoint was saved at step: {ckpt['step']}\")\n",
     "\n",
     "# Load model weights\n",
-    "# Load into unwrapped model (checkpoint saved without DataParallel 'module.' prefix)\n",
-    "load_target = model.module if hasattr(model, 'module') else model\n",
-    "load_target.load_state_dict(ckpt['model_state_dict'])\n",
+    "# Load into unwrapped model (model_unwrapped set in cell 10)\n",
+    "model_unwrapped.load_state_dict(ckpt['model_state_dict'])\n",
     "print(\"Model weights loaded\")\n",
     "\n",
     "# Load EMA weights\n",
@@ -900,7 +899,7 @@
     "resume_step = ckpt['step']\n",
     "if 'optimizer_state_dict' in ckpt:\n",
     "    optimizer = torch.optim.AdamW(\n",
-    "        model.parameters(),\n",
+    "        model_unwrapped.parameters(),\n",
     "        lr=config.learning_rate,\n",
     "        betas=(0.9, 0.98),\n",
     "        weight_decay=config.weight_decay,\n",