diff --git a/lab2/PT_Part1_MNIST.ipynb b/lab2/PT_Part1_MNIST.ipynb
index bcfae490..74e0e076 100644
--- a/lab2/PT_Part1_MNIST.ipynb
+++ b/lab2/PT_Part1_MNIST.ipynb
@@ -413,7 +413,7 @@
         "        correct_pred = 0\n",
         "        total_pred = 0\n",
         "\n",
-        "        for images, labels in trainset_loader:\n",
+        "        for images, labels in dataloader:\n",
         "            # Move tensors to GPU so compatible with model\n",
         "            images, labels = images.to(device), labels.to(device)\n",
         "\n",
@@ -498,7 +498,7 @@
         "    total_pred = 0\n",
         "    # Disable gradient calculations when in inference mode\n",
         "    with torch.no_grad():\n",
-        "        for images, labels in testset_loader:\n",
+        "        for images, labels in dataloader:\n",
         "            # TODO: ensure evalaution happens on the GPU\n",
         "            images, labels = # TODO\n",
         "\n",