diff --git a/lab2/PT_Part1_MNIST.ipynb b/lab2/PT_Part1_MNIST.ipynb
index 4a345b27..7bcde853 100644
--- a/lab2/PT_Part1_MNIST.ipynb
+++ b/lab2/PT_Part1_MNIST.ipynb
@@ -418,12 +418,12 @@
         "            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",
diff --git a/lab2/solutions/PT_Part1_MNIST_Solution.ipynb b/lab2/solutions/PT_Part1_MNIST_Solution.ipynb
index efb23831..779a6011 100644
--- a/lab2/solutions/PT_Part1_MNIST_Solution.ipynb
+++ b/lab2/solutions/PT_Part1_MNIST_Solution.ipynb
@@ -423,9 +423,9 @@
         "            # 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",