Upload main.ipynb with huggingface_hub

main.ipynb

@@ -915,6 +915,13 @@
     "del ckpt  # Free memory\n",
     "torch.cuda.empty_cache()\n",
     "\n",
+    "# Set up DataParallel if multiple GPUs available\n",
+    "if torch.cuda.device_count() > 1:\n",
+    "    model_dp = nn.DataParallel(model_unwrapped, device_ids=[0, 1], output_device=0)\n",
+    "    print(f\"\\nUsing {torch.cuda.device_count()} GPUs with DataParallel!\")\n",
+    "else:\n",
+    "    model_dp = model_unwrapped\n",
+    "\n",
     "print(f\"\\nReady to resume training from step {resume_step}\")\n",
     "print(f\"Or skip to generation cells to use the model!\")\n"
    ]