diff --git a/lab2/solutions/PT_Part1_MNIST_Solution.ipynb b/lab2/solutions/PT_Part1_MNIST_Solution.ipynb
index efb23831..c1423d96 100644
--- a/lab2/solutions/PT_Part1_MNIST_Solution.ipynb
+++ b/lab2/solutions/PT_Part1_MNIST_Solution.ipynb
@@ -915,6 +915,9 @@
         "# Process test set in batches\n",
         "with torch.no_grad():\n",
         "    for images, labels in testset_loader:\n",
+        "        # TODO: ensure evalaution happens on the GPU\n",
+        "        images, labels = images.to(device), labels.to(device)\n",
+        "\n",
         "        outputs = cnn_model(images)\n",
         "\n",
         "        # Apply softmax to get probabilities from the predicted logits\n",