Update notebook too

Author: adamjstewart

examples/camvid_segmentation_multiclass.ipynb

@@ -492,7 +492,7 @@
     "    augmentation=get_validation_augmentation(),\n",
     ")\n",
     "\n",
-    "#Change to > 0 if not on Windows machine\n",
+    "# Change to > 0 if not on Windows machine\n",
     "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=0)\n",
     "valid_loader = DataLoader(valid_dataset, batch_size=32, shuffle=False, num_workers=0)\n",
     "test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=0)"
@@ -545,12 +545,10 @@
     "import pytorch_lightning as pl\n",
     "import segmentation_models_pytorch as smp\n",
     "import torch\n",
-    "import torch.nn.functional as F\n",
     "from torch.optim import lr_scheduler\n",
     "\n",
     "\n",
     "class CamVidModel(pl.LightningModule):\n",
-    "\n",
     "    def __init__(self, arch, encoder_name, in_channels, out_classes, **kwargs):\n",
     "        super().__init__()\n",
     "        self.model = smp.create_model(\n",
@@ -591,13 +589,14 @@
     "        mask = mask.long()\n",
     "\n",
     "        # Mask shape\n",
-    "        assert mask.ndim == 3 # [batch_size, H, W]\n",
+    "        assert mask.ndim == 3  # [batch_size, H, W]\n",
     "\n",
     "        # Predict mask logits\n",
     "        logits_mask = self.forward(image)\n",
-    "        \n",
-    "        assert logits_mask.shape[1] == self.number_of_classes # [batch_size, number_of_classes, H, W]\n",
-    "        \n",
+    "\n",
+    "        assert (\n",
+    "            logits_mask.shape[1] == self.number_of_classes\n",
+    "        )  # [batch_size, number_of_classes, H, W]\n",
     "\n",
     "        # Ensure the logits mask is contiguous\n",
     "        logits_mask = logits_mask.contiguous()\n",
@@ -1678,7 +1677,6 @@
     }
    ],
    "source": [
-    "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
     "# Fetch a batch from the test loader\n",