fix: pointnet example (#449)

wandb · Nov 3, 2023 · 0ea7668 · 0ea7668
1 parent 3f52406
commit 0ea7668
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/colabs/pyg/pointnet-classification/02_pointnet_plus_plus.ipynb b/colabs/pyg/pointnet-classification/02_pointnet_plus_plus.ipynb
@@ -364,8 +364,9 @@
     "        range(num_train_examples),\n",
     "        desc=f\"Training Epoch {epoch}/{config.epochs}\"\n",
     "    )\n",
+    "    data_iter = iter(train_loader)\n",
     "    for batch_idx in progress_bar:\n",
-    "        data = next(iter(train_loader)).to(device)\n",
+    "        data = next(data_iter).to(device)\n",
     "        \n",
     "        optimizer.zero_grad()\n",
     "        prediction = model(data)\n",
@@ -395,8 +396,9 @@
     "        range(num_val_examples),\n",
     "        desc=f\"Validation Epoch {epoch}/{config.epochs}\"\n",
     "    )\n",
+    "    data_iter = iter(val_loader)\n",
     "    for batch_idx in progress_bar:\n",
-    "        data = next(iter(val_loader)).to(device)\n",
+    "        data = next(data_iter).to(device)\n",
     "        \n",
     "        with torch.no_grad():\n",
     "            prediction = model(data)\n",

diff --git a/colabs/pyg/pointnet-classification/03_sweep.ipynb b/colabs/pyg/pointnet-classification/03_sweep.ipynb
@@ -286,8 +286,9 @@
     "            range(num_train_examples),\n",
     "            desc=f\"Training Epoch {epoch}/{config.epochs}\"\n",
     "        )\n",
+    "        data_iter = iter(train_loader)\n",
     "        for batch_idx in progress_bar:\n",
-    "            data = next(iter(train_loader)).to(device)\n",
+    "            data = next(data_iter).to(device)\n",
     "\n",
     "            optimizer.zero_grad()\n",
     "            prediction = model(data)\n",
@@ -315,8 +316,9 @@
     "            range(num_val_examples),\n",
     "            desc=f\"Validation Epoch {epoch}/{config.epochs}\"\n",
     "        )\n",
+    "        data_iter = iter(val_loader)\n",
     "        for batch_idx in progress_bar:\n",
-    "            data = next(iter(val_loader)).to(device)\n",
+    "            data = next(data_iter).to(device)\n",
     "\n",
     "            with torch.no_grad():\n",
     "                prediction = model(data)\n",