diff --git a/course/training/01_core.ipynb b/course/training/01_core.ipynb
index 5dc4366..d7fc3c3 100644
--- a/course/training/01_core.ipynb
+++ b/course/training/01_core.ipynb
@@ -326,7 +326,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "softmax = torch.nn.functional.softmax(outputs.prediction_logits[0], dim=0)  # create probability distribution\n",
+    "softmax = torch.nn.functional.softmax(outputs.prediction_logits[0], dim=1)  # create probability distribution\n",
     "argmax = torch.argmax(softmax, dim=1)  # get index of the max probability"
    ]
   },