diff --git a/lab2/PT_Part1_MNIST.ipynb b/lab2/PT_Part1_MNIST.ipynb
index bcfae490..9f48fb51 100644
--- a/lab2/PT_Part1_MNIST.ipynb
+++ b/lab2/PT_Part1_MNIST.ipynb
@@ -413,17 +413,17 @@
         "        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",
         "            # Forward pass\n",
-        "            outputs = fc_model(images)\n",
+        "            outputs = model(images)\n",
         "\n",
         "            # Clear gradients before performing backward pass\n",
         "            optimizer.zero_grad()\n",
         "            # Calculate loss based on model predictions\n",
-        "            loss = loss_function(outputs, labels)\n",
+        "            loss = criterion(outputs, labels)\n",
         "            # Backpropagate and update model parameters\n",
         "            loss.backward()\n",
         "            optimizer.step()\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",
diff --git a/lab2/solutions/PT_Part1_MNIST_Solution.ipynb b/lab2/solutions/PT_Part1_MNIST_Solution.ipynb
index 01dc5bfa..5ab77df6 100644
--- a/lab2/solutions/PT_Part1_MNIST_Solution.ipynb
+++ b/lab2/solutions/PT_Part1_MNIST_Solution.ipynb
@@ -417,15 +417,15 @@
         "        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",
         "            # Clear gradients before performing backward pass\n",
         "            optimizer.zero_grad()\n",
         "            # Forward pass\n",
-        "            outputs = fc_model(images)\n",
+        "            outputs = model(images)\n",
         "            # Calculate loss based on model predictions\n",
-        "            loss = loss_function(outputs, labels)\n",
+        "            loss = criterion(outputs, labels)\n",
         "            # Backpropagate and update model parameters\n",
         "            loss.backward()\n",
         "            optimizer.step()\n",
@@ -500,7 +500,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 = images.to(device), labels.to(device)\n",
         "            # images, labels = # TODO\n",
@@ -533,7 +533,7 @@
         "    return test_loss, test_acc\n",
         "\n",
         "# TODO: call the evaluate function to evaluate the trained model!!\n",
-        "test_loss, test_acc = evaluate(fc_model, trainset_loader, loss_function)\n",
+        "test_loss, test_acc = evaluate(fc_model, testset_loader, loss_function)\n",
         "# test_loss, test_acc = # TODO\n",
         "\n",
         "print('Test accuracy:', test_acc)"
@@ -1026,4 +1026,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
\ No newline at end of file
+}