use torch.argmax instead of torch.max to remove tuple unpacking confusion

Author: vyinkabanjo

lab2/solutions/PT_Part1_MNIST_Solution.ipynb

@@ -428,7 +428,7 @@
         "            total_loss += loss.item()*images.size(0)\n",
         "\n",
         "            # Calculate accuracy\n",
-        "            _, predicted = torch.max(outputs, 1)  # Get predicted class\n",
+        "            predicted = torch.argmax(outputs, dim=1)  # Get predicted class\n",
         "            correct_pred += (predicted == labels).sum().item()  # Count correct predictions\n",
         "            total_pred += labels.size(0) # Count total predictions\n",
         "\n",
@@ -509,7 +509,7 @@
         "            # test_loss += TODO\n",
         "\n",
         "            # Calculate accuracy\n",
-        "            _, predicted = torch.max(outputs, 1)\n",
+        "            predicted = torch.argmax(outputs, dim=1)\n",
         "            #TODO: identify the digit with the highest confidence prediction for the first image in the test dataset.\n",
         "\n",
         "            correct_pred += (predicted == labels).sum().item()\n",
@@ -736,7 +736,7 @@
         "        optimizer.step()\n",
         "\n",
         "        # Calculate accuracy\n",
-        "        _, predicted = torch.max(logits, 1) \n",
+        "        predicted = torch.argmax(logits, dim=1)\n",
         "        correct_pred += (predicted == labels).sum().item() \n",
         "        total_pred += labels.size(0)\n",
         "\n",
@@ -907,7 +907,7 @@
         "        probabilities = torch.nn.functional.softmax(outputs, dim=1)\n",
         "\n",
         "        # Get predicted classes\n",
-        "        _, predicted = torch.max(outputs, 1)\n",
+        "        predicted = torch.argmax(outputs, dim=1)\n",
         "\n",
         "        all_predictions.append(probabilities) \n",
         "        all_labels.append(labels)\n",