diff --git "a/ViTFINAL.ipynb" "b/ViTFINAL.ipynb"
new file mode 100644--- /dev/null
+++ "b/ViTFINAL.ipynb"
@@ -0,0 +1,5319 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "b00dbc42",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[XPU] Cache cleared and peak memory stats reset.\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "\n",
+    "def clean_device_memory():\n",
+    "    if torch.xpu.is_available():\n",
+    "        torch.xpu.empty_cache()\n",
+    "        torch.xpu.reset_peak_memory_stats()\n",
+    "        print(\"[XPU] Cache cleared and peak memory stats reset.\")\n",
+    "    elif torch.cuda.is_available():\n",
+    "        torch.cuda.empty_cache()\n",
+    "        torch.cuda.reset_peak_memory_stats()\n",
+    "        print(\"[CUDA] Cache cleared and peak memory stats reset.\")\n",
+    "    else:\n",
+    "        print(\"No XPU or CUDA device found.\")\n",
+    "\n",
+    "if __name__ == '__main__':\n",
+    "    clean_device_memory()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "598b91f6",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "c:\\Python310\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "--- Starting ViT Training with Larger Batch Size ---\n",
+      "Using device: xpu with batch size: 16\n",
+      "Torch version: 2.7.0+xpu\n",
+      "Current XPU memory allocated: 0.0 MB\n",
+      "Max XPU memory allocated: 0.0 MB\n",
+      "\n",
+      "[1/4] Loading dataset...\n",
+      "[DEBUG] Found 10676 real images\n",
+      "[DEBUG] Found 76182 fake images\n",
+      "[DEBUG] Using 10676 samples per class (balanced)\n",
+      "[DEBUG] Final dataset size: 21352 samples\n",
+      "[DEBUG] Batch size: 16 | Total batches: 1335\n",
+      "\n",
+      "[2/4] Initializing model...\n",
+      "\n",
+      "[Memory] Allocated: 327.30 MB | Reserved: 382.00 MB | Max: 327.30 MB\n",
+      "\n",
+      "[3/4] Starting training...\n",
+      "\n",
+      "[Epoch 1/10]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   0%|          | 0/1335 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 327.30 MB | Reserved: 382.00 MB | Max: 327.30 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   4%|▎         | 50/1335 [01:05<26:31,  1.24s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.77 MB | Reserved: 1592.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   7%|▋         | 100/1335 [02:06<25:17,  1.23s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.01 MB | Reserved: 1624.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  11%|█         | 150/1335 [03:08<24:17,  1.23s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.77 MB | Reserved: 1622.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  15%|█▍        | 200/1335 [04:09<23:06,  1.22s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.49 MB | Reserved: 1612.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  19%|█▊        | 250/1335 [05:10<21:54,  1.21s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.77 MB | Reserved: 1602.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  22%|██▏       | 300/1335 [06:11<21:05,  1.22s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.48 MB | Reserved: 1584.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  26%|██▌       | 350/1335 [07:12<20:06,  1.22s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.30 MB | Reserved: 1612.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  30%|██▉       | 400/1335 [08:13<18:59,  1.22s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.77 MB | Reserved: 1622.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  34%|███▎      | 450/1335 [09:14<17:57,  1.22s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.49 MB | Reserved: 1614.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  37%|███▋      | 500/1335 [10:15<16:59,  1.22s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.71 MB | Reserved: 1592.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  41%|████      | 550/1335 [11:16<15:48,  1.21s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.77 MB | Reserved: 1602.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  45%|████▍     | 600/1335 [12:17<15:00,  1.22s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.27 MB | Reserved: 1592.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  49%|████▊     | 650/1335 [13:18<13:52,  1.22s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.01 MB | Reserved: 1624.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  52%|█████▏    | 700/1335 [14:19<12:59,  1.23s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.77 MB | Reserved: 1622.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  56%|█████▌    | 750/1335 [15:20<11:57,  1.23s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.77 MB | Reserved: 1602.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  60%|█████▉    | 800/1335 [16:22<11:01,  1.24s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.77 MB | Reserved: 1622.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  64%|██████▎   | 850/1335 [17:25<10:23,  1.29s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.49 MB | Reserved: 1614.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  67%|██████▋   | 900/1335 [18:28<09:15,  1.28s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.27 MB | Reserved: 1592.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  71%|███████   | 950/1335 [19:31<08:08,  1.27s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.77 MB | Reserved: 1622.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  75%|███████▍  | 1000/1335 [20:35<07:10,  1.28s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.48 MB | Reserved: 1584.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  79%|███████▊  | 1050/1335 [21:38<06:01,  1.27s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.26 MB | Reserved: 1614.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  82%|████████▏ | 1100/1335 [22:42<04:59,  1.28s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.24 MB | Reserved: 1620.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  86%|████████▌ | 1150/1335 [23:45<03:52,  1.26s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.27 MB | Reserved: 1592.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  90%|████████▉ | 1200/1335 [24:49<02:53,  1.28s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.48 MB | Reserved: 1584.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  94%|█████████▎| 1250/1335 [25:52<01:47,  1.27s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.24 MB | Reserved: 1620.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  97%|█████████▋| 1300/1335 [26:56<00:44,  1.26s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.27 MB | Reserved: 1592.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training: 100%|██████████| 1335/1335 [27:40<00:00,  1.24s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[Results] Loss: 923.3108 | Accuracy: 54.88%\n",
+      "\n",
+      "[Memory] Allocated: 1331.43 MB | Reserved: 1640.00 MB | Max: 3180.14 MB\n",
+      "\n",
+      "[Epoch 2/10]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   0%|          | 0/1335 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.18 MB | Reserved: 1640.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   4%|▎         | 50/1335 [01:00<25:35,  1.19s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   7%|▋         | 100/1335 [02:00<24:17,  1.18s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  11%|█         | 150/1335 [02:59<23:21,  1.18s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  15%|█▍        | 200/1335 [03:57<21:53,  1.16s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  19%|█▊        | 250/1335 [04:55<20:51,  1.15s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  22%|██▏       | 300/1335 [05:54<20:04,  1.16s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  26%|██▌       | 350/1335 [06:52<18:54,  1.15s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  30%|██▉       | 400/1335 [07:48<17:22,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  34%|███▎      | 450/1335 [08:44<16:14,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  37%|███▋      | 500/1335 [09:39<15:34,  1.12s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  41%|████      | 550/1335 [10:34<14:21,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  45%|████▍     | 600/1335 [11:30<13:28,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  49%|████▊     | 650/1335 [12:25<12:34,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  52%|█████▏    | 700/1335 [13:20<11:41,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  56%|█████▌    | 750/1335 [14:15<10:44,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  60%|█████▉    | 800/1335 [15:10<09:49,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  64%|██████▎   | 850/1335 [16:06<08:56,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  67%|██████▋   | 900/1335 [17:01<07:57,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  71%|███████   | 950/1335 [17:56<07:04,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  75%|███████▍  | 1000/1335 [18:51<06:08,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  79%|███████▊  | 1050/1335 [19:46<05:14,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  82%|████████▏ | 1100/1335 [20:41<04:18,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  86%|████████▌ | 1150/1335 [21:36<03:23,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  90%|████████▉ | 1200/1335 [22:32<02:27,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  94%|█████████▎| 1250/1335 [23:27<01:33,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  97%|█████████▋| 1300/1335 [24:23<00:38,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training: 100%|██████████| 1335/1335 [25:01<00:00,  1.12s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[Results] Loss: 875.8758 | Accuracy: 60.55%\n",
+      "\n",
+      "[Memory] Allocated: 1331.26 MB | Reserved: 1612.00 MB | Max: 3180.14 MB\n",
+      "\n",
+      "[Epoch 3/10]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   0%|          | 0/1335 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.01 MB | Reserved: 1612.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   4%|▎         | 50/1335 [00:55<23:37,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   7%|▋         | 100/1335 [01:50<22:36,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  11%|█         | 150/1335 [02:46<21:58,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  15%|█▍        | 200/1335 [03:41<20:50,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  19%|█▊        | 250/1335 [04:36<19:54,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  22%|██▏       | 300/1335 [05:31<19:17,  1.12s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  26%|██▌       | 350/1335 [06:26<17:55,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  30%|██▉       | 400/1335 [07:21<17:03,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  34%|███▎      | 450/1335 [08:15<16:00,  1.08s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  37%|███▋      | 500/1335 [09:11<16:06,  1.16s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  41%|████      | 550/1335 [10:05<14:26,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  45%|████▍     | 600/1335 [11:00<13:22,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  49%|████▊     | 650/1335 [11:55<12:26,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  52%|█████▏    | 700/1335 [12:49<11:33,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  56%|█████▌    | 750/1335 [13:44<10:37,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  60%|█████▉    | 800/1335 [14:39<09:44,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  64%|██████▎   | 850/1335 [15:34<08:51,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  67%|██████▋   | 900/1335 [16:28<07:56,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  71%|███████   | 950/1335 [17:23<07:00,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  75%|███████▍  | 1000/1335 [18:18<06:06,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  79%|███████▊  | 1050/1335 [19:12<05:12,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  82%|████████▏ | 1100/1335 [20:07<04:17,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  86%|████████▌ | 1150/1335 [21:02<03:21,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  90%|████████▉ | 1200/1335 [21:57<02:27,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  94%|█████████▎| 1250/1335 [22:51<01:33,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  97%|█████████▋| 1300/1335 [23:46<00:38,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training: 100%|██████████| 1335/1335 [24:24<00:00,  1.10s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[Results] Loss: 846.1695 | Accuracy: 63.43%\n",
+      "\n",
+      "[Memory] Allocated: 1331.50 MB | Reserved: 1598.00 MB | Max: 3180.14 MB\n",
+      "\n",
+      "[Epoch 4/10]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   0%|          | 0/1335 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.00 MB | Reserved: 1598.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   4%|▎         | 50/1335 [00:54<23:22,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.43 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   7%|▋         | 100/1335 [01:49<22:30,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1324.91 MB | Reserved: 1646.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  11%|█         | 150/1335 [02:44<21:30,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.38 MB | Reserved: 1646.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  15%|█▍        | 200/1335 [03:38<20:35,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.43 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  19%|█▊        | 250/1335 [04:33<19:43,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.38 MB | Reserved: 1646.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  22%|██▏       | 300/1335 [05:27<18:45,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.43 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  26%|██▌       | 350/1335 [06:22<17:51,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.38 MB | Reserved: 1646.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  30%|██▉       | 400/1335 [07:17<17:02,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.43 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  34%|███▎      | 450/1335 [08:11<16:03,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.43 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  37%|███▋      | 500/1335 [09:06<15:08,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.38 MB | Reserved: 1646.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  41%|████      | 550/1335 [10:01<14:14,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.43 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  45%|████▍     | 600/1335 [10:55<13:20,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1324.91 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  49%|████▊     | 650/1335 [11:50<12:28,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.38 MB | Reserved: 1646.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  52%|█████▏    | 700/1335 [12:44<11:32,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.38 MB | Reserved: 1646.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  56%|█████▌    | 750/1335 [13:39<10:41,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.43 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  60%|█████▉    | 800/1335 [14:33<09:42,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1324.91 MB | Reserved: 1646.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  64%|██████▎   | 850/1335 [15:28<08:49,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1324.91 MB | Reserved: 1646.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  67%|██████▋   | 900/1335 [16:23<07:54,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.43 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  71%|███████   | 950/1335 [17:17<06:58,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.43 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  75%|███████▍  | 1000/1335 [18:12<06:07,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.43 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  79%|███████▊  | 1050/1335 [19:06<05:11,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.43 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  82%|████████▏ | 1100/1335 [20:01<04:16,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.43 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  86%|████████▌ | 1150/1335 [20:55<03:21,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.43 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  90%|████████▉ | 1200/1335 [21:50<02:28,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1324.91 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  94%|█████████▎| 1250/1335 [22:45<01:33,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.38 MB | Reserved: 1646.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  97%|█████████▋| 1300/1335 [23:39<00:38,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1324.91 MB | Reserved: 1636.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training: 100%|██████████| 1335/1335 [24:17<00:00,  1.09s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[Results] Loss: 810.5119 | Accuracy: 66.20%\n",
+      "\n",
+      "[Memory] Allocated: 1331.22 MB | Reserved: 1620.00 MB | Max: 3180.14 MB\n",
+      "\n",
+      "[Epoch 5/10]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   0%|          | 0/1335 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.97 MB | Reserved: 1620.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   4%|▎         | 50/1335 [00:55<23:41,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   7%|▋         | 100/1335 [01:50<22:40,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  11%|█         | 150/1335 [02:45<21:50,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  15%|█▍        | 200/1335 [03:40<20:55,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  19%|█▊        | 250/1335 [04:36<20:05,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  22%|██▏       | 300/1335 [05:31<19:00,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  26%|██▌       | 350/1335 [06:27<18:23,  1.12s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  30%|██▉       | 400/1335 [07:22<17:14,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  34%|███▎      | 450/1335 [08:17<16:20,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  37%|███▋      | 500/1335 [09:13<15:34,  1.12s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  41%|████      | 550/1335 [10:08<14:21,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  45%|████▍     | 600/1335 [11:04<13:33,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  49%|████▊     | 650/1335 [11:59<12:38,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  52%|█████▏    | 700/1335 [12:54<11:46,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  56%|█████▌    | 750/1335 [13:50<10:44,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  60%|█████▉    | 800/1335 [14:45<09:52,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  64%|██████▎   | 850/1335 [15:40<08:52,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  67%|██████▋   | 900/1335 [16:35<07:57,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  71%|███████   | 950/1335 [17:30<07:05,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  75%|███████▍  | 1000/1335 [18:26<06:12,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  79%|███████▊  | 1050/1335 [19:21<05:16,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  82%|████████▏ | 1100/1335 [20:16<04:20,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  86%|████████▌ | 1150/1335 [21:12<03:25,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  90%|████████▉ | 1200/1335 [22:07<02:29,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  94%|█████████▎| 1250/1335 [23:02<01:35,  1.12s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  97%|█████████▋| 1300/1335 [23:57<00:38,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training: 100%|██████████| 1335/1335 [24:36<00:00,  1.11s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[Results] Loss: 777.2801 | Accuracy: 68.05%\n",
+      "\n",
+      "[Memory] Allocated: 1331.26 MB | Reserved: 1612.00 MB | Max: 3180.14 MB\n",
+      "\n",
+      "[Epoch 6/10]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   0%|          | 0/1335 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.01 MB | Reserved: 1612.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   4%|▎         | 50/1335 [00:55<23:45,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   7%|▋         | 100/1335 [01:50<22:44,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  11%|█         | 150/1335 [02:46<21:53,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  15%|█▍        | 200/1335 [03:41<20:48,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  19%|█▊        | 250/1335 [04:37<20:19,  1.12s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  22%|██▏       | 300/1335 [05:32<19:04,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  26%|██▌       | 350/1335 [06:27<18:07,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  30%|██▉       | 400/1335 [07:23<17:23,  1.12s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  34%|███▎      | 450/1335 [08:18<16:18,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  37%|███▋      | 500/1335 [09:14<15:29,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  41%|████      | 550/1335 [10:09<14:31,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  45%|████▍     | 600/1335 [11:04<13:28,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  49%|████▊     | 650/1335 [12:00<12:40,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  52%|█████▏    | 700/1335 [12:55<11:43,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  56%|█████▌    | 750/1335 [13:51<10:51,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  60%|█████▉    | 800/1335 [14:46<09:50,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  64%|██████▎   | 850/1335 [15:42<08:57,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  67%|██████▋   | 900/1335 [16:37<07:58,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  71%|███████   | 950/1335 [17:32<07:09,  1.12s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  75%|███████▍  | 1000/1335 [18:28<06:11,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  79%|███████▊  | 1050/1335 [19:23<05:16,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  82%|████████▏ | 1100/1335 [20:19<04:21,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  86%|████████▌ | 1150/1335 [21:14<03:23,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  90%|████████▉ | 1200/1335 [22:10<02:29,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  94%|█████████▎| 1250/1335 [23:05<01:34,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  97%|█████████▋| 1300/1335 [24:00<00:38,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training: 100%|██████████| 1335/1335 [24:39<00:00,  1.11s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[Results] Loss: 738.1488 | Accuracy: 70.40%\n",
+      "\n",
+      "[Memory] Allocated: 1331.50 MB | Reserved: 1598.00 MB | Max: 3180.14 MB\n",
+      "\n",
+      "[Epoch 7/10]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   0%|          | 0/1335 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.00 MB | Reserved: 1598.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   4%|▎         | 50/1335 [00:55<23:42,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   7%|▋         | 100/1335 [01:51<22:51,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  11%|█         | 150/1335 [02:46<21:45,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  15%|█▍        | 200/1335 [03:41<21:14,  1.12s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1324.68 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  19%|█▊        | 250/1335 [04:37<20:00,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  22%|██▏       | 300/1335 [05:32<19:00,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  26%|██▌       | 350/1335 [06:28<18:21,  1.12s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  30%|██▉       | 400/1335 [07:23<17:12,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  34%|███▎      | 450/1335 [08:19<16:14,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  37%|███▋      | 500/1335 [09:14<15:18,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  41%|████      | 550/1335 [10:10<14:27,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  45%|████▍     | 600/1335 [11:05<13:28,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  49%|████▊     | 650/1335 [12:01<12:39,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  52%|█████▏    | 700/1335 [12:56<11:41,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  56%|█████▌    | 750/1335 [13:52<10:47,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  60%|█████▉    | 800/1335 [14:47<09:51,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  64%|██████▎   | 850/1335 [15:42<09:02,  1.12s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  67%|██████▋   | 900/1335 [16:38<08:01,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  71%|███████   | 950/1335 [17:33<07:06,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  75%|███████▍  | 1000/1335 [18:29<06:10,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  79%|███████▊  | 1050/1335 [19:24<05:17,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  82%|████████▏ | 1100/1335 [20:20<04:20,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  86%|████████▌ | 1150/1335 [21:15<03:24,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  90%|████████▉ | 1200/1335 [22:11<02:30,  1.12s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  94%|█████████▎| 1250/1335 [23:06<01:34,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  97%|█████████▋| 1300/1335 [24:02<00:39,  1.12s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training: 100%|██████████| 1335/1335 [24:40<00:00,  1.11s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[Results] Loss: 703.8742 | Accuracy: 72.20%\n",
+      "\n",
+      "[Memory] Allocated: 1331.26 MB | Reserved: 1612.00 MB | Max: 3180.14 MB\n",
+      "\n",
+      "[Epoch 8/10]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   0%|          | 0/1335 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.01 MB | Reserved: 1612.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   4%|▎         | 50/1335 [00:55<23:45,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   7%|▋         | 100/1335 [01:51<22:47,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  11%|█         | 150/1335 [02:46<21:36,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  15%|█▍        | 200/1335 [03:42<20:44,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  19%|█▊        | 250/1335 [04:36<19:44,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  22%|██▏       | 300/1335 [05:31<18:51,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  26%|██▌       | 350/1335 [06:26<17:56,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  30%|██▉       | 400/1335 [07:20<17:16,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  34%|███▎      | 450/1335 [08:15<16:10,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  37%|███▋      | 500/1335 [09:10<15:13,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  41%|████      | 550/1335 [10:05<14:17,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  45%|████▍     | 600/1335 [10:59<13:24,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  49%|████▊     | 650/1335 [11:54<12:30,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  52%|█████▏    | 700/1335 [12:49<11:35,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  56%|█████▌    | 750/1335 [13:44<10:38,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  60%|█████▉    | 800/1335 [14:38<09:48,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  64%|██████▎   | 850/1335 [15:33<08:55,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  67%|██████▋   | 900/1335 [16:28<07:57,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  71%|███████   | 950/1335 [17:23<07:03,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  75%|███████▍  | 1000/1335 [18:18<06:08,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  79%|███████▊  | 1050/1335 [19:13<05:12,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  82%|████████▏ | 1100/1335 [20:07<04:16,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  86%|████████▌ | 1150/1335 [21:02<03:22,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  90%|████████▉ | 1200/1335 [21:57<02:27,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  94%|█████████▎| 1250/1335 [22:52<01:32,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  97%|█████████▋| 1300/1335 [23:47<00:38,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training: 100%|██████████| 1335/1335 [24:24<00:00,  1.10s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[Results] Loss: 667.9493 | Accuracy: 74.06%\n",
+      "\n",
+      "[Memory] Allocated: 1331.50 MB | Reserved: 1598.00 MB | Max: 3180.14 MB\n",
+      "\n",
+      "[Epoch 9/10]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   0%|          | 0/1335 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.00 MB | Reserved: 1598.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   4%|▎         | 50/1335 [00:54<23:30,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   7%|▋         | 100/1335 [01:49<22:36,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  11%|█         | 150/1335 [02:44<21:40,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  15%|█▍        | 200/1335 [03:39<20:37,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  19%|█▊        | 250/1335 [04:34<19:51,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  22%|██▏       | 300/1335 [05:28<19:01,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  26%|██▌       | 350/1335 [06:23<17:59,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  30%|██▉       | 400/1335 [07:18<17:10,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  34%|███▎      | 450/1335 [08:13<16:20,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  37%|███▋      | 500/1335 [09:08<15:13,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  41%|████      | 550/1335 [10:03<14:20,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  45%|████▍     | 600/1335 [10:58<13:24,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  49%|████▊     | 650/1335 [11:52<12:32,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  52%|█████▏    | 700/1335 [12:47<11:39,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  56%|█████▌    | 750/1335 [13:42<10:39,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  60%|█████▉    | 800/1335 [14:37<09:50,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  64%|██████▎   | 850/1335 [15:32<08:51,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  67%|██████▋   | 900/1335 [16:27<07:55,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  71%|███████   | 950/1335 [17:22<07:01,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  75%|█████��█▍  | 1000/1335 [18:16<06:09,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  79%|███████▊  | 1050/1335 [19:11<05:12,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  82%|████████▏ | 1100/1335 [20:06<04:16,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  86%|████████▌ | 1150/1335 [21:01<03:23,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  90%|████████▉ | 1200/1335 [21:56<02:28,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  94%|█████████▎| 1250/1335 [22:51<01:34,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  97%|█████████▋| 1300/1335 [23:46<00:38,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.19 MB | Reserved: 1668.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training: 100%|██████████| 1335/1335 [24:24<00:00,  1.10s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[Results] Loss: 641.1695 | Accuracy: 75.52%\n",
+      "\n",
+      "[Memory] Allocated: 1331.26 MB | Reserved: 1612.00 MB | Max: 3180.14 MB\n",
+      "\n",
+      "[Epoch 10/10]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   0%|          | 0/1335 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1326.01 MB | Reserved: 1612.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   4%|▎         | 50/1335 [00:54<23:22,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:   7%|▋         | 100/1335 [01:49<22:33,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  11%|█         | 150/1335 [02:44<21:48,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  15%|█▍        | 200/1335 [03:39<20:45,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  19%|█▊        | 250/1335 [04:34<19:54,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  22%|██▏       | 300/1335 [05:29<18:59,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  26%|██▌       | 350/1335 [06:24<18:03,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  30%|██▉       | 400/1335 [07:19<17:08,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  34%|███▎      | 450/1335 [08:14<16:10,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  37%|███▋      | 500/1335 [09:09<15:17,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  41%|████      | 550/1335 [10:04<14:28,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  45%|████▍     | 600/1335 [10:59<13:28,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  49%|████▊     | 650/1335 [11:54<12:31,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  52%|█████▏    | 700/1335 [12:49<11:39,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  56%|█████▌    | 750/1335 [13:44<10:44,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  60%|█████▉    | 800/1335 [14:39<09:47,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  64%|██████▎   | 850/1335 [15:34<08:49,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  67%|██████▋   | 900/1335 [16:28<07:59,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  71%|███████   | 950/1335 [17:23<07:02,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  75%|███████▍  | 1000/1335 [18:18<06:07,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  79%|███████▊  | 1050/1335 [19:13<05:12,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  82%|████████▏ | 1100/1335 [20:08<04:16,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  86%|████████▌ | 1150/1335 [21:03<03:23,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  90%|████████▉ | 1200/1335 [21:58<02:28,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  94%|█████████▎| 1250/1335 [22:52<01:32,  1.09s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training:  97%|█████████▋| 1300/1335 [23:47<00:38,  1.10s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "[Memory] Allocated: 1325.46 MB | Reserved: 1670.00 MB | Max: 3180.14 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training: 100%|██████████| 1335/1335 [24:25<00:00,  1.10s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[Results] Loss: 603.1762 | Accuracy: 77.53%\n",
+      "\n",
+      "[Memory] Allocated: 1331.50 MB | Reserved: 1598.00 MB | Max: 3180.14 MB\n",
+      "\n",
+      "[4/4] Saving model...\n",
+      "Model saved to ./vit_digital_forensics_xpu.pth\n"
+     ]
+    }
+   ],
+   "source": [
+    "# ViT Training with Larger Batch Size on XPU (Fixed Dataset Class)\n",
+    "\n",
+    "import os\n",
+    "import glob\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from torchvision import transforms\n",
+    "from PIL import Image\n",
+    "import timm\n",
+    "from tqdm import tqdm\n",
+    "import random\n",
+    "\n",
+    "# ----------------------------\n",
+    "# CONFIGURATION\n",
+    "# ----------------------------\n",
+    "DATASET_DIR = 'frames'\n",
+    "BATCH_SIZE = 16  # Increased from 8 to 16\n",
+    "EPOCHS = 10\n",
+    "LEARNING_RATE = 3e-5\n",
+    "NUM_CLASSES = 2\n",
+    "IMG_SIZE = 224\n",
+    "DEVICE = 'xpu' if torch.xpu.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'\n",
+    "MODEL_SAVE_PATH = './vit_digital_forensics_xpu.pth'\n",
+    "\n",
+    "print(\"\\n--- Starting ViT Training with Larger Batch Size ---\")\n",
+    "print(f\"Using device: {DEVICE} with batch size: {BATCH_SIZE}\")\n",
+    "print(\"Torch version:\", torch.__version__)\n",
+    "print(\"Current XPU memory allocated:\", torch.xpu.memory_allocated()/1024**2, \"MB\")\n",
+    "print(\"Max XPU memory allocated:\", torch.xpu.max_memory_allocated()/1024**2, \"MB\")\n",
+    "\n",
+    "# ----------------------------\n",
+    "# DATASET CLASS (FIXED IMPLEMENTATION)\n",
+    "# ----------------------------\n",
+    "class ForgeryImageDataset(Dataset):\n",
+    "    def __init__(self, root_dir, transform=None):\n",
+    "        self.samples = []\n",
+    "        self.transform = transform\n",
+    "        \n",
+    "        # Verify dataset structure\n",
+    "        real_folder = os.path.join(root_dir, 'real')\n",
+    "        fake_folder = os.path.join(root_dir, 'fake')\n",
+    "        print(f\"[DEBUG] Found {len(glob.glob(os.path.join(real_folder, '*.jpg')))} real images\")\n",
+    "        print(f\"[DEBUG] Found {len(glob.glob(os.path.join(fake_folder, '*.jpg')))} fake images\")\n",
+    "        \n",
+    "        for label, name in enumerate(['real', 'fake']):\n",
+    "            folder = os.path.join(root_dir, name)\n",
+    "            for img_path in glob.glob(os.path.join(folder, '*.jpg')):\n",
+    "                self.samples.append((img_path, label))\n",
+    "\n",
+    "        # Balance classes\n",
+    "        real_samples = [s for s in self.samples if s[1] == 0]\n",
+    "        fake_samples = [s for s in self.samples if s[1] == 1]\n",
+    "        min_len = min(len(real_samples), len(fake_samples))\n",
+    "        print(f\"[DEBUG] Using {min_len} samples per class (balanced)\")\n",
+    "        \n",
+    "        self.samples = random.sample(real_samples, min_len) + random.sample(fake_samples, min_len)\n",
+    "        random.shuffle(self.samples)\n",
+    "        print(f\"[DEBUG] Final dataset size: {len(self.samples)} samples\")\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.samples)\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        img_path, label = self.samples[idx]\n",
+    "        image = Image.open(img_path).convert('RGB')\n",
+    "        if self.transform:\n",
+    "            image = self.transform(image)\n",
+    "        return image, label\n",
+    "\n",
+    "# ----------------------------\n",
+    "# TRANSFORMS\n",
+    "# ----------------------------\n",
+    "transform = transforms.Compose([\n",
+    "    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n",
+    "    transforms.RandomHorizontalFlip(),\n",
+    "    transforms.RandomRotation(10),\n",
+    "    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize([0.5]*3, [0.5]*3)\n",
+    "])\n",
+    "\n",
+    "# ----------------------------\n",
+    "# MEMORY MONITORING FUNCTION\n",
+    "# ----------------------------\n",
+    "def print_memory_stats():\n",
+    "    if DEVICE == 'xpu':\n",
+    "        stats = {\n",
+    "            'allocated': torch.xpu.memory_allocated()/1024**2,\n",
+    "            'reserved': torch.xpu.memory_reserved()/1024**2,\n",
+    "            'max_alloc': torch.xpu.max_memory_allocated()/1024**2\n",
+    "        }\n",
+    "        print(\"\\n[Memory] Allocated: {allocated:.2f} MB | Reserved: {reserved:.2f} MB | Max: {max_alloc:.2f} MB\".format(**stats))\n",
+    "\n",
+    "# ----------------------------\n",
+    "# TRAINING LOOP\n",
+    "# ----------------------------\n",
+    "def train():\n",
+    "    # Initialize dataset and loader\n",
+    "    print(\"\\n[1/4] Loading dataset...\")\n",
+    "    dataset = ForgeryImageDataset(DATASET_DIR, transform=transform)\n",
+    "    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)\n",
+    "    print(f\"[DEBUG] Batch size: {BATCH_SIZE} | Total batches: {len(loader)}\")\n",
+    "    \n",
+    "    # Initialize model\n",
+    "    print(\"\\n[2/4] Initializing model...\")\n",
+    "    model = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=NUM_CLASSES)\n",
+    "    model = model.to(DEVICE)\n",
+    "    criterion = nn.CrossEntropyLoss()\n",
+    "    optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE)\n",
+    "    print_memory_stats()\n",
+    "    \n",
+    "    # Training loop\n",
+    "    print(\"\\n[3/4] Starting training...\")\n",
+    "    model.train()\n",
+    "    for epoch in range(EPOCHS):\n",
+    "        print(f\"\\n[Epoch {epoch+1}/{EPOCHS}]\")\n",
+    "        total_loss, correct = 0, 0\n",
+    "        \n",
+    "        for batch_idx, (imgs, labels) in enumerate(tqdm(loader, desc=\"Training\")):\n",
+    "            # Memory check every 50 batches\n",
+    "            if batch_idx % 50 == 0:\n",
+    "                print_memory_stats()\n",
+    "            \n",
+    "            # Move data to device\n",
+    "            imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)\n",
+    "            \n",
+    "            # Forward pass\n",
+    "            optimizer.zero_grad()\n",
+    "            outputs = model(imgs)\n",
+    "            loss = criterion(outputs, labels)\n",
+    "            \n",
+    "            # Backward pass\n",
+    "            loss.backward()\n",
+    "            optimizer.step()\n",
+    "            \n",
+    "            # Metrics\n",
+    "            total_loss += loss.item()\n",
+    "            correct += (torch.argmax(outputs, dim=1) == labels).sum().item()\n",
+    "            \n",
+    "            # Clean up\n",
+    "            del outputs\n",
+    "            torch.xpu.empty_cache()\n",
+    "        \n",
+    "        # Epoch statistics\n",
+    "        epoch_acc = 100. * correct / len(dataset)\n",
+    "        print(f\"[Results] Loss: {total_loss:.4f} | Accuracy: {epoch_acc:.2f}%\")\n",
+    "        print_memory_stats()\n",
+    "\n",
+    "    # Save model\n",
+    "    print(\"\\n[4/4] Saving model...\")\n",
+    "    torch.save(model.state_dict(), MODEL_SAVE_PATH)\n",
+    "    print(f\"Model saved to {MODEL_SAVE_PATH}\")\n",
+    "\n",
+    "if __name__ == '__main__':\n",
+    "    train()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "7459c841",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "c:\\Python310\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n",
+      "[I 2025-05-18 12:22:12,743] A new study created in memory with name: no-name-0238caf2-2544-421d-8841-7ed656b2c736\n",
+      "C:\\Users\\devcloud\\AppData\\Local\\Temp\\ipykernel_12100\\463889153.py:98: FutureWarning: suggest_loguniform has been deprecated in v3.0.0. This feature will be removed in v6.0.0. See https://github.com/optuna/optuna/releases/tag/v3.0.0. Use suggest_float(..., log=True) instead.\n",
+      "  lr = trial.suggest_loguniform('lr', 1e-5, 1e-3)\n",
+      "C:\\Users\\devcloud\\AppData\\Local\\Temp\\ipykernel_12100\\463889153.py:100: FutureWarning: suggest_loguniform has been deprecated in v3.0.0. This feature will be removed in v6.0.0. See https://github.com/optuna/optuna/releases/tag/v3.0.0. Use suggest_float(..., log=True) instead.\n",
+      "  weight_decay = trial.suggest_loguniform('weight_decay', 1e-6, 1e-2)\n",
+      "[I 2025-05-18 13:36:53,991] Trial 0 finished with value: 0.5346604215456674 and parameters: {'lr': 0.00019442949702926413, 'batch_size': 8, 'weight_decay': 0.0066764067597068905}. Best is trial 0 with value: 0.5346604215456674.\n",
+      "[I 2025-05-18 14:44:15,574] Trial 1 finished with value: 0.5381733021077283 and parameters: {'lr': 0.00010325233112691653, 'batch_size': 16, 'weight_decay': 1.5499675871668855e-05}. Best is trial 1 with value: 0.5381733021077283.\n",
+      "[I 2025-05-18 15:54:16,517] Trial 2 finished with value: 0.7782201405152225 and parameters: {'lr': 1.0369533463569177e-05, 'batch_size': 32, 'weight_decay': 0.009006790996778456}. Best is trial 2 with value: 0.7782201405152225.\n",
+      "[I 2025-05-18 17:03:08,338] Trial 3 finished with value: 0.5091334894613583 and parameters: {'lr': 0.0003046327737509729, 'batch_size': 16, 'weight_decay': 7.688366508752892e-06}. Best is trial 2 with value: 0.7782201405152225.\n",
+      "[I 2025-05-18 18:11:08,096] Trial 4 finished with value: 0.5110070257611241 and parameters: {'lr': 0.0008134874694963872, 'batch_size': 32, 'weight_decay': 0.0012103868091215206}. Best is trial 2 with value: 0.7782201405152225.\n",
+      "[I 2025-05-18 19:20:58,079] Trial 5 finished with value: 0.7440281030444965 and parameters: {'lr': 1.0562531428879552e-05, 'batch_size': 32, 'weight_decay': 0.0021515636255839536}. Best is trial 2 with value: 0.7782201405152225.\n",
+      "[I 2025-05-18 20:31:06,385] Trial 6 finished with value: 0.5014051522248244 and parameters: {'lr': 0.0006120287586805153, 'batch_size': 16, 'weight_decay': 0.0003485698669576425}. Best is trial 2 with value: 0.7782201405152225.\n",
+      "[I 2025-05-18 21:38:58,408] Trial 7 finished with value: 0.6770491803278689 and parameters: {'lr': 3.2343860924914236e-05, 'batch_size': 16, 'weight_decay': 1.1758959449886972e-05}. Best is trial 2 with value: 0.7782201405152225.\n",
+      "[I 2025-05-18 22:56:43,453] Trial 8 finished with value: 0.4960187353629977 and parameters: {'lr': 0.0009598946035424287, 'batch_size': 8, 'weight_decay': 0.007537464093474474}. Best is trial 2 with value: 0.7782201405152225.\n",
+      "[I 2025-05-19 00:13:39,096] Trial 9 finished with value: 0.6517564402810304 and parameters: {'lr': 2.5743923801217405e-05, 'batch_size': 8, 'weight_decay': 1.3856930026507902e-06}. Best is trial 2 with value: 0.7782201405152225.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Best hyperparameters found:\n",
+      "{'lr': 1.0369533463569177e-05, 'batch_size': 32, 'weight_decay': 0.009006790996778456}\n",
+      "Accuracy: 77.82%\n",
+      "Model saved at: vit_best_model_optuna.pth\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import glob\n",
+    "import random\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from torch.utils.data import Dataset, DataLoader, random_split\n",
+    "from torchvision import transforms\n",
+    "from PIL import Image\n",
+    "import timm\n",
+    "from tqdm import tqdm\n",
+    "import optuna\n",
+    "\n",
+    "# ----------------------------\n",
+    "# CONFIGURATION\n",
+    "# ----------------------------\n",
+    "DATASET_DIR = \"frames\"\n",
+    "DEVICE = 'xpu' if torch.xpu.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'\n",
+    "IMG_SIZE = 224\n",
+    "NUM_CLASSES = 2\n",
+    "MODEL_SAVE_PATH = \"vit_best_model_optuna.pth\"\n",
+    "\n",
+    "# ----------------------------\n",
+    "# DATASET CLASS\n",
+    "# ----------------------------\n",
+    "class ForgeryImageDataset(Dataset):\n",
+    "    def __init__(self, root_dir, transform=None):\n",
+    "        self.samples = []\n",
+    "        self.transform = transform\n",
+    "\n",
+    "        for label, name in enumerate(['real', 'fake']):\n",
+    "            folder = os.path.join(root_dir, name)\n",
+    "            for img_path in glob.glob(os.path.join(folder, '*.jpg')):\n",
+    "                self.samples.append((img_path, label))\n",
+    "\n",
+    "        # Balance classes\n",
+    "        real_samples = [s for s in self.samples if s[1] == 0]\n",
+    "        fake_samples = [s for s in self.samples if s[1] == 1]\n",
+    "        min_len = min(len(real_samples), len(fake_samples))\n",
+    "        self.samples = random.sample(real_samples, min_len) + random.sample(fake_samples, min_len)\n",
+    "        random.shuffle(self.samples)\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.samples)\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        img_path, label = self.samples[idx]\n",
+    "        image = Image.open(img_path).convert('RGB')\n",
+    "        if self.transform:\n",
+    "            image = self.transform(image)\n",
+    "        return image, label\n",
+    "\n",
+    "# ----------------------------\n",
+    "# TRANSFORMS\n",
+    "# ----------------------------\n",
+    "transform = transforms.Compose([\n",
+    "    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n",
+    "    transforms.RandomHorizontalFlip(),\n",
+    "    transforms.ColorJitter(0.1, 0.1, 0.1),\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize([0.5]*3, [0.5]*3)\n",
+    "])\n",
+    "\n",
+    "# ----------------------------\n",
+    "# TRAINING FUNCTION\n",
+    "# ----------------------------\n",
+    "def train_model(model, train_loader, optimizer, criterion):\n",
+    "    model.train()\n",
+    "    for imgs, labels in train_loader:\n",
+    "        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)\n",
+    "        optimizer.zero_grad()\n",
+    "        outputs = model(imgs)\n",
+    "        loss = criterion(outputs, labels)\n",
+    "        loss.backward()\n",
+    "        optimizer.step()\n",
+    "\n",
+    "# ----------------------------\n",
+    "# VALIDATION FUNCTION\n",
+    "# ----------------------------\n",
+    "def evaluate_model(model, val_loader):\n",
+    "    model.eval()\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "    with torch.no_grad():\n",
+    "        for imgs, labels in val_loader:\n",
+    "            imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)\n",
+    "            outputs = model(imgs)\n",
+    "            preds = outputs.argmax(1)\n",
+    "            correct += (preds == labels).sum().item()\n",
+    "            total += labels.size(0)\n",
+    "    return correct / total\n",
+    "\n",
+    "# ----------------------------\n",
+    "# OPTUNA OBJECTIVE FUNCTION\n",
+    "# ----------------------------\n",
+    "def objective(trial):\n",
+    "    # Hyperparameters to tune\n",
+    "    lr = trial.suggest_loguniform('lr', 1e-5, 1e-3)\n",
+    "    batch_size = trial.suggest_categorical('batch_size', [8, 16, 32])\n",
+    "    weight_decay = trial.suggest_loguniform('weight_decay', 1e-6, 1e-2)\n",
+    "\n",
+    "    # Prepare dataset\n",
+    "    dataset = ForgeryImageDataset(DATASET_DIR, transform=transform)\n",
+    "    val_len = int(0.2 * len(dataset))\n",
+    "    train_len = len(dataset) - val_len\n",
+    "    train_ds, val_ds = random_split(dataset, [train_len, val_len])\n",
+    "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)\n",
+    "    val_loader = DataLoader(val_ds, batch_size=batch_size)\n",
+    "\n",
+    "    # Model\n",
+    "    model = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=NUM_CLASSES)\n",
+    "    model = model.to(DEVICE)\n",
+    "\n",
+    "    criterion = nn.CrossEntropyLoss()\n",
+    "    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
+    "\n",
+    "    # Training loop (fewer epochs to save time)\n",
+    "    EPOCHS = 5\n",
+    "    for _ in range(EPOCHS):\n",
+    "        train_model(model, train_loader, optimizer, criterion)\n",
+    "\n",
+    "    acc = evaluate_model(model, val_loader)\n",
+    "    trial.set_user_attr(\"model\", model.state_dict())\n",
+    "    return acc\n",
+    "\n",
+    "# ----------------------------\n",
+    "# RUN OPTUNA\n",
+    "# ----------------------------\n",
+    "study = optuna.create_study(direction=\"maximize\")\n",
+    "study.optimize(objective, n_trials=10)\n",
+    "\n",
+    "# ----------------------------\n",
+    "# SAVE BEST MODEL\n",
+    "# ----------------------------\n",
+    "best_trial = study.best_trial\n",
+    "best_model_state = best_trial.user_attrs[\"model\"]\n",
+    "\n",
+    "final_model = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=NUM_CLASSES)\n",
+    "final_model.load_state_dict(best_model_state)\n",
+    "torch.save(final_model.state_dict(), MODEL_SAVE_PATH)\n",
+    "\n",
+    "print(\"\\nBest hyperparameters found:\")\n",
+    "print(best_trial.params)\n",
+    "print(f\"Accuracy: {best_trial.value:.2%}\")\n",
+    "print(f\"Model saved at: {MODEL_SAVE_PATH}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "764639bf",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "c:\\Python310\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using device: xpu\n",
+      "\n",
+      "===== Training vit_base_patch16_224 =====\n",
+      "Epoch 1/5 | Train Acc: 68.58% | Val Acc: 76.79%\n",
+      "Epoch 2/5 | Train Acc: 81.67% | Val Acc: 80.18%\n",
+      "Epoch 3/5 | Train Acc: 85.97% | Val Acc: 86.66%\n",
+      "Epoch 4/5 | Train Acc: 88.99% | Val Acc: 85.64%\n",
+      "Epoch 5/5 | Train Acc: 90.65% | Val Acc: 85.36%\n",
+      "\n",
+      "===== Training vit_base_patch32_224 =====\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "c:\\Python310\\lib\\site-packages\\huggingface_hub\\file_download.py:143: UserWarning: `huggingface_hub` cache-system uses symlinks by default to efficiently store duplicated files but your machine does not support them in C:\\Users\\devcloud\\.cache\\huggingface\\hub\\models--timm--vit_base_patch32_224.augreg_in21k_ft_in1k. Caching files will still work but in a degraded version that might require more space on your disk. This warning can be disabled by setting the `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For more details, see https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations.\n",
+      "To support symlinks on Windows, you either need to activate Developer Mode or to run Python as an administrator. In order to activate developer mode, see this article: https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development\n",
+      "  warnings.warn(message)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/5 | Train Acc: 61.57% | Val Acc: 71.96%\n",
+      "Epoch 2/5 | Train Acc: 69.05% | Val Acc: 74.73%\n",
+      "Epoch 3/5 | Train Acc: 73.89% | Val Acc: 74.59%\n",
+      "Epoch 4/5 | Train Acc: 76.91% | Val Acc: 77.67%\n",
+      "Epoch 5/5 | Train Acc: 79.54% | Val Acc: 79.54%\n",
+      "\n",
+      "===== Training vit_small_patch16_224 =====\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "c:\\Python310\\lib\\site-packages\\huggingface_hub\\file_download.py:143: UserWarning: `huggingface_hub` cache-system uses symlinks by default to efficiently store duplicated files but your machine does not support them in C:\\Users\\devcloud\\.cache\\huggingface\\hub\\models--timm--vit_small_patch16_224.augreg_in21k_ft_in1k. Caching files will still work but in a degraded version that might require more space on your disk. This warning can be disabled by setting the `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For more details, see https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations.\n",
+      "To support symlinks on Windows, you either need to activate Developer Mode or to run Python as an administrator. In order to activate developer mode, see this article: https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development\n",
+      "  warnings.warn(message)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/5 | Train Acc: 64.53% | Val Acc: 71.56%\n",
+      "Epoch 2/5 | Train Acc: 75.12% | Val Acc: 83.55%\n",
+      "Epoch 3/5 | Train Acc: 80.67% | Val Acc: 81.25%\n",
+      "Epoch 4/5 | Train Acc: 83.72% | Val Acc: 85.68%\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import time\n",
+    "import csv\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from torch.utils.data import DataLoader, WeightedRandomSampler\n",
+    "from torchvision import transforms, datasets\n",
+    "import timm\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "# ===================== CONFIG =====================\n",
+    "DATA_DIR = \"frames\"  # root that contains train/ and val/\n",
+    "OUTPUT_DIR = os.path.join(DATA_DIR, \"model_outputs\")\n",
+    "CSV_PATH = os.path.join(OUTPUT_DIR, \"vit_comparison_results.csv\")\n",
+    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
+    "\n",
+    "DEVICE = (\n",
+    "    torch.device(\"xpu\") if torch.xpu.is_available()\n",
+    "    else torch.device(\"cuda\") if torch.cuda.is_available()\n",
+    "    else torch.device(\"cpu\")\n",
+    ")\n",
+    "print(\"Using device:\", DEVICE)\n",
+    "\n",
+    "# Hyper‑parameters fixed from Optuna best trial\n",
+    "EPOCHS = 5\n",
+    "BATCH_SIZE = 32\n",
+    "LR = 1.0369533463569177e-05\n",
+    "WEIGHT_DECAY = 0.009006790996778456\n",
+    "IMG_SIZE = 224\n",
+    "NUM_CLASSES = 2\n",
+    "\n",
+    "# Vision Transformer variants to compare (timm names)\n",
+    "MODEL_NAMES = [\n",
+    "    \"vit_large_patch16_224\",\n",
+    "    \"deit_base_patch16_224\",\n",
+    "    \"swin_tiny_patch4_window7_224\",\n",
+    "]\n",
+    "\n",
+    "# ===================== DATA =====================\n",
+    "transform_train = transforms.Compose([\n",
+    "    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n",
+    "    transforms.RandomHorizontalFlip(),\n",
+    "    transforms.RandomRotation(10),\n",
+    "    transforms.ColorJitter(0.1, 0.1, 0.1),\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize([0.5]*3, [0.5]*3),\n",
+    "])\n",
+    "\n",
+    "transform_val = transforms.Compose([\n",
+    "    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize([0.5]*3, [0.5]*3),\n",
+    "])\n",
+    "\n",
+    "train_ds = datasets.ImageFolder(os.path.join(DATA_DIR, \"train\"), transform=transform_train)\n",
+    "val_ds   = datasets.ImageFolder(os.path.join(DATA_DIR, \"val\"),   transform=transform_val)\n",
+    "\n",
+    "# ----- Balanced Sampler for train loader -----\n",
+    "class_counts = [0]*NUM_CLASSES\n",
+    "for _ , label in train_ds.samples:\n",
+    "    class_counts[label] += 1\n",
+    "class_weights = [sum(class_counts)/c for c in class_counts]\n",
+    "sample_weights = [class_weights[label] for _ , label in train_ds.samples]\n",
+    "train_sampler = WeightedRandomSampler(sample_weights, num_samples=len(train_ds), replacement=True)\n",
+    "\n",
+    "train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, sampler=train_sampler, num_workers=0)\n",
+    "val_loader   = DataLoader(val_ds,   batch_size=BATCH_SIZE, shuffle=False,   num_workers=0)\n",
+    "\n",
+    "# ===================== TRAIN & EVAL FUNCTIONS =====================\n",
+    "\n",
+    "def train_one_epoch(model, loader, criterion, optimizer):\n",
+    "    model.train()\n",
+    "    running_loss, running_correct, total = 0.0, 0, 0\n",
+    "    for imgs, labels in loader:\n",
+    "        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)\n",
+    "        optimizer.zero_grad()\n",
+    "        outputs = model(imgs)\n",
+    "        loss = criterion(outputs, labels)\n",
+    "        loss.backward()\n",
+    "        optimizer.step()\n",
+    "\n",
+    "        running_loss += loss.item() * imgs.size(0)\n",
+    "        running_correct += (outputs.argmax(1) == labels).sum().item()\n",
+    "        total += labels.size(0)\n",
+    "    return running_loss/total, running_correct/total\n",
+    "\n",
+    "\n",
+    "def evaluate(model, loader, criterion):\n",
+    "    model.eval()\n",
+    "    loss_sum, correct, total = 0.0, 0, 0\n",
+    "    with torch.no_grad():\n",
+    "        for imgs, labels in loader:\n",
+    "            imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)\n",
+    "            outputs = model(imgs)\n",
+    "            loss = criterion(outputs, labels)\n",
+    "            loss_sum += loss.item()*imgs.size(0)\n",
+    "            correct += (outputs.argmax(1)==labels).sum().item()\n",
+    "            total += labels.size(0)\n",
+    "    return loss_sum/total, correct/total\n",
+    "\n",
+    "# ===================== MAIN LOOP =====================\n",
+    "results = []\n",
+    "criterion = nn.CrossEntropyLoss()\n",
+    "\n",
+    "for model_name in MODEL_NAMES:\n",
+    "    print(f\"\\n===== Training {model_name} =====\")\n",
+    "    model = timm.create_model(model_name, pretrained=True, num_classes=NUM_CLASSES)\n",
+    "    model.to(DEVICE)\n",
+    "\n",
+    "    optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n",
+    "\n",
+    "    start_time = time.time()\n",
+    "    for epoch in range(1, EPOCHS+1):\n",
+    "        train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer)\n",
+    "        val_loss, val_acc     = evaluate(model, val_loader, criterion)\n",
+    "        print(f\"Epoch {epoch}/{EPOCHS} | Train Acc: {train_acc*100:.2f}% | Val Acc: {val_acc*100:.2f}%\")\n",
+    "\n",
+    "    elapsed = (time.time() - start_time)/60  # minutes\n",
+    "    model_path = os.path.join(OUTPUT_DIR, f\"{model_name.replace('/', '_')}.pth\")\n",
+    "    torch.save(model.state_dict(), model_path)\n",
+    "\n",
+    "    results.append({\n",
+    "        'model': model_name,\n",
+    "        'val_accuracy': f\"{val_acc*100:.2f}\",\n",
+    "        'train_accuracy': f\"{train_acc*100:.2f}\",\n",
+    "        'time_min': f\"{elapsed:.1f}\",\n",
+    "        'weights_path': model_path\n",
+    "    })\n",
+    "\n",
+    "# ===================== SAVE CSV =====================\n",
+    "with open(CSV_PATH, 'w', newline='') as csvfile:\n",
+    "    writer = csv.DictWriter(csvfile, fieldnames=results[0].keys())\n",
+    "    writer.writeheader()\n",
+    "    for row in results:\n",
+    "        writer.writerow(row)\n",
+    "\n",
+    "print(\"\\n🏁 All models trained. Results saved to:\", CSV_PATH)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "cd8a5e3c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "c:\\Python310\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using device: xpu\n",
+      "\n",
+      "===== Training vit_small_patch32_224 =====\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                         \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/5 | Train Acc: 87.94% | Val Acc: 88.12%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                         \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 2/5 | Train Acc: 88.36% | Val Acc: 88.20%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                         \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 3/5 | Train Acc: 89.15% | Val Acc: 88.60%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                         \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 4/5 | Train Acc: 90.41% | Val Acc: 88.95%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                         \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 5/5 | Train Acc: 91.58% | Val Acc: 88.38%\n",
+      "Model saved to model_outputs\\vit_small_patch32_224.pth\n",
+      "\n",
+      "===== Training vit_tiny_patch16_224 =====\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                         \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/5 | Train Acc: 87.85% | Val Acc: 87.91%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                         \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 2/5 | Train Acc: 88.47% | Val Acc: 88.59%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                         \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 3/5 | Train Acc: 89.47% | Val Acc: 88.69%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                         \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 4/5 | Train Acc: 90.95% | Val Acc: 88.95%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                         \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 5/5 | Train Acc: 92.41% | Val Acc: 89.20%\n",
+      "Model saved to model_outputs\\vit_tiny_patch16_224.pth\n",
+      "\n",
+      "===== Training vit_tiny_patch32_224 =====\n",
+      "Error training vit_tiny_patch32_224: Unknown model (vit_tiny_patch32_224)\n",
+      "\n",
+      "===== Training vit_large_patch16_224 =====\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "c:\\Python310\\lib\\site-packages\\huggingface_hub\\file_download.py:143: UserWarning: `huggingface_hub` cache-system uses symlinks by default to efficiently store duplicated files but your machine does not support them in C:\\Users\\devcloud\\.cache\\huggingface\\hub\\models--timm--vit_large_patch16_224.augreg_in21k_ft_in1k. Caching files will still work but in a degraded version that might require more space on your disk. This warning can be disabled by setting the `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For more details, see https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations.\n",
+      "To support symlinks on Windows, you either need to activate Developer Mode or to run Python as an administrator. In order to activate developer mode, see this article: https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development\n",
+      "  warnings.warn(message)\n",
+      "                                                                             \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/5 | Train Acc: 88.04% | Val Acc: 88.32%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                             \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 2/5 | Train Acc: 89.53% | Val Acc: 89.89%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                             \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 3/5 | Train Acc: 91.60% | Val Acc: 90.58%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                             \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 4/5 | Train Acc: 93.46% | Val Acc: 90.49%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                             \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 5/5 | Train Acc: 94.89% | Val Acc: 91.22%\n",
+      "Model saved to model_outputs\\vit_large_patch16_224.pth\n",
+      "\n",
+      "===== Training vit_large_patch32_224 =====\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "c:\\Python310\\lib\\site-packages\\huggingface_hub\\file_download.py:143: UserWarning: `huggingface_hub` cache-system uses symlinks by default to efficiently store duplicated files but your machine does not support them in C:\\Users\\devcloud\\.cache\\huggingface\\hub\\models--timm--vit_large_patch32_224.orig_in21k. Caching files will still work but in a degraded version that might require more space on your disk. This warning can be disabled by setting the `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For more details, see https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations.\n",
+      "To support symlinks on Windows, you either need to activate Developer Mode or to run Python as an administrator. In order to activate developer mode, see this article: https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development\n",
+      "  warnings.warn(message)\n",
+      "                                                                           \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/5 | Train Acc: 87.97% | Val Acc: 88.16%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                           \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 2/5 | Train Acc: 88.31% | Val Acc: 88.42%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                           \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 3/5 | Train Acc: 89.06% | Val Acc: 88.70%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                           \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 4/5 | Train Acc: 90.13% | Val Acc: 89.34%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                           \r"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 5/5 | Train Acc: 91.31% | Val Acc: 89.23%\n",
+      "Model saved to model_outputs\\vit_large_patch32_224.pth\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import torch\n",
+    "from torch import nn\n",
+    "from torchvision import datasets, transforms\n",
+    "from torch.utils.data import DataLoader\n",
+    "from timm import create_model\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "# Paths\n",
+    "data_dir = \"frames\"\n",
+    "\n",
+    "# Device\n",
+    "device = torch.device(\"xpu\" if torch.xpu.is_available() else \"cpu\")\n",
+    "print(f\"Using device: {device}\")\n",
+    "\n",
+    "# Hyperparameters\n",
+    "BATCH_SIZE = 32\n",
+    "LR = 1.03695e-5\n",
+    "WEIGHT_DECAY = 0.0090067\n",
+    "EPOCHS = 5\n",
+    "\n",
+    "# Transform\n",
+    "transform = transforms.Compose([\n",
+    "    transforms.Resize((224, 224)),\n",
+    "    transforms.ToTensor(),\n",
+    "])\n",
+    "\n",
+    "# Datasets & Dataloaders\n",
+    "train_dataset = datasets.ImageFolder(os.path.join(data_dir, 'train'), transform=transform)\n",
+    "val_dataset = datasets.ImageFolder(os.path.join(data_dir, 'val'), transform=transform)\n",
+    "\n",
+    "train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
+    "val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)\n",
+    "\n",
+    "# Vision Transformer model variants\n",
+    "vit_models = [\n",
+    "    \"vit_small_patch32_224\",\n",
+    "    \"vit_tiny_patch16_224\",\n",
+    "    \"vit_tiny_patch32_224\",\n",
+    "    \"vit_large_patch16_224\",\n",
+    "    \"vit_large_patch32_224\",\n",
+    "    # \"vit_huge_patch14_224\"  # Uncomment if you have enough resources\n",
+    "]\n",
+    "\n",
+    "# Training Function\n",
+    "def train_model(model_name):\n",
+    "    print(f\"\\n===== Training {model_name} =====\")\n",
+    "\n",
+    "    model = create_model(model_name, pretrained=True, num_classes=2)\n",
+    "    model.to(device)\n",
+    "\n",
+    "    criterion = nn.CrossEntropyLoss()\n",
+    "    optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n",
+    "\n",
+    "    for epoch in range(EPOCHS):\n",
+    "        model.train()\n",
+    "        train_correct = 0\n",
+    "        total_train = 0\n",
+    "\n",
+    "        for images, labels in tqdm(train_loader, desc=f\"Epoch {epoch+1}/{EPOCHS} - Training\", leave=False):\n",
+    "            images, labels = images.to(device), labels.to(device)\n",
+    "            optimizer.zero_grad()\n",
+    "            outputs = model(images)\n",
+    "            loss = criterion(outputs, labels)\n",
+    "            loss.backward()\n",
+    "            optimizer.step()\n",
+    "\n",
+    "            _, preds = torch.max(outputs, 1)\n",
+    "            train_correct += (preds == labels).sum().item()\n",
+    "            total_train += labels.size(0)\n",
+    "\n",
+    "        train_acc = 100 * train_correct / total_train\n",
+    "\n",
+    "        # Evaluation\n",
+    "        model.eval()\n",
+    "        val_correct = 0\n",
+    "        total_val = 0\n",
+    "        with torch.no_grad():\n",
+    "            for images, labels in val_loader:\n",
+    "                images, labels = images.to(device), labels.to(device)\n",
+    "                outputs = model(images)\n",
+    "                _, preds = torch.max(outputs, 1)\n",
+    "                val_correct += (preds == labels).sum().item()\n",
+    "                total_val += labels.size(0)\n",
+    "\n",
+    "        val_acc = 100 * val_correct / total_val\n",
+    "        print(f\"Epoch {epoch+1}/{EPOCHS} | Train Acc: {train_acc:.2f}% | Val Acc: {val_acc:.2f}%\")\n",
+    "\n",
+    "    # Save model weights\n",
+    "    os.makedirs(\"model_outputs\", exist_ok=True)\n",
+    "    save_path = os.path.join(\"model_outputs\", f\"{model_name}.pth\")\n",
+    "    torch.save(model.state_dict(), save_path)\n",
+    "    print(f\"Model saved to {save_path}\")\n",
+    "\n",
+    "# Run training for each ViT model\n",
+    "for model_name in vit_models:\n",
+    "    try:\n",
+    "        train_model(model_name)\n",
+    "    except Exception as e:\n",
+    "        print(f\"Error training {model_name}: {str(e)}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}