smaller model and better tensorboard instructions

01_CARE/solution.ipynb

@@ -660,7 +660,7 @@
    "outputs": [],
    "source": [
     "# Load the model\n",
-    "model = UNet(depth=3, in_channels=1, out_channels=1)"
+    "model = UNet(depth=2, in_channels=1, out_channels=1)"
    ]
   },
   {
@@ -791,10 +791,9 @@
     "![image](nb_data/extensions.png)\n",
     "\n",
     "2) Search Tensorboard and install and install the extension published by Microsoft.\n",
-    "3) Open the command palette (ctrl+shift+p), search for Python: Launch Tensorboard and hit enter.\n",
-    "4) When prompted, select either \"Use current working directory\" or \"Select another folder\" and enter the path to the `01_CARE/runs/` directory.\n",
-    "\n",
-    "Once it's open, continue to the next cell to start training, then open the Tensorboard tab to see the loss curves.\n",
+    "3) Start training. Run the cell below to begin training the model and generating logs.\n",
+    "3) Once training is started. Open the command palette (ctrl+shift+p), search for Python: Launch Tensorboard and hit enter.\n",
+    "4) When prompted, select \"Select another folder\" and enter the path to the `01_CARE/runs/` directory.\n",
     "\n",
     "</div>"
    ]
@@ -817,7 +816,7 @@
    "outputs": [],
    "source": [
     "# Training loop\n",
-    "n_epochs = 10\n",
+    "n_epochs = 5\n",
     "device = \"cuda\" if cuda.is_available() else \"cpu\"\n",
     "model.to(device)\n",
     "\n",