diff --git a/tutorials/W1D2_ComparingTasks/W1D2_Tutorial1.ipynb b/tutorials/W1D2_ComparingTasks/W1D2_Tutorial1.ipynb
index b9f5b6b8d..8e022bb4f 100644
--- a/tutorials/W1D2_ComparingTasks/W1D2_Tutorial1.ipynb
+++ b/tutorials/W1D2_ComparingTasks/W1D2_Tutorial1.ipynb
@@ -176,6 +176,8 @@
     "from pathlib import Path\n",
     "import zipfile\n",
     "import random\n",
+    "import contextlib\n",
+    "import io\n",
     "\n",
     "# Import third-party libraries\n",
     "import numpy as np\n",
@@ -844,50 +846,49 @@
    },
    "outputs": [],
    "source": [
-    "# Define a transformation pipeline for the MNIST dataset\n",
-    "mnist_transform = transforms.Compose([\n",
-    "    transforms.Resize((32, 32)),  # Resize the images to 32x32 pixels\n",
-    "    transforms.ToTensor(),  # Convert images to PyTorch tensors\n",
-    "    transforms.Normalize(mean=(0.1307,), std=(0.3081,))  # Normalize the images with mean and standard deviation\n",
-    "])\n",
+    "with contextlib.redirect_stdout(io.StringIO()):\n",
     "\n",
-    "# Load the MNIST training dataset with transformations applied\n",
-    "train_val_dataset = torchvision.datasets.MNIST(\n",
-    "    root='./data',  # Directory to store/load the data\n",
-    "    train=True,  # Specify to load the training set\n",
-    "    transform=mnist_transform,  # Apply the transformation pipeline defined earlier\n",
-    "    download=True  # Download the dataset if it's not already present\n",
-    ")\n",
+    "    # Define a transformation pipeline for the MNIST dataset\n",
+    "    mnist_transform = transforms.Compose([\n",
+    "        transforms.Resize((32, 32)),  # Resize the images to 32x32 pixels\n",
+    "        transforms.ToTensor(),  # Convert images to PyTorch tensors\n",
+    "        transforms.Normalize(mean=(0.1307,), std=(0.3081,))  # Normalize the images with mean and standard deviation\n",
+    "    ])\n",
     "\n",
-    "# Load the MNIST test dataset with transformations applied\n",
-    "test_dataset = torchvision.datasets.MNIST(\n",
-    "    root='./data',  # Directory to store/load the data\n",
-    "    train=False,  # Specify to load the test set\n",
-    "    transform=mnist_transform,  # Apply the transformation pipeline defined earlier\n",
-    "    download=True  # Download the dataset if it's not already present\n",
-    ")\n",
+    "    # Load the MNIST training dataset with transformations applied\n",
+    "    train_val_dataset = torchvision.datasets.MNIST(\n",
+    "        root='./data',  # Directory to store/load the data\n",
+    "        train=True,  # Specify to load the training set\n",
+    "        transform=mnist_transform,  # Apply the transformation pipeline defined earlier\n",
+    "        download=True  # Download the dataset if it's not already present\n",
+    "    )\n",
     "\n",
-    "# Split the training dataset into training and validation sets\n",
-    "train_size = int(0.9 * len(train_val_dataset))  # Calculate the size of the training set (90% of the original)\n",
-    "val_size = len(train_val_dataset) - train_size  # Calculate the size of the validation set (remaining 10%)\n",
-    "train_dataset, val_dataset = torch.utils.data.random_split(\n",
-    "    dataset=train_val_dataset,  # Original training dataset to split\n",
-    "    lengths=[train_size, val_size]  # Lengths of the resulting splits\n",
-    ")\n",
+    "    # Load the MNIST test dataset with transformations applied\n",
+    "    test_dataset = torchvision.datasets.MNIST(\n",
+    "        root='./data',  # Directory to store/load the data\n",
+    "        train=False,  # Specify to load the test set\n",
+    "        transform=mnist_transform,  # Apply the transformation pipeline defined earlier\n",
+    "        download=True  # Download the dataset if it's not already present\n",
+    "    )\n",
     "\n",
-    "# Split the test dataset into two halves: original and transfer sets\n",
-    "test_size_original = int(0.5 * len(test_dataset))  # Calculate the size of the original test set (50% of the original)\n",
-    "test_size_transfer = len(test_dataset) - test_size_original  # Calculate the size of the transfer test set (remaining 50%)\n",
-    "test_dataset_original, test_dataset_transfer = torch.utils.data.random_split(\n",
-    "    dataset=test_dataset,  # Original test dataset to split\n",
-    "    lengths=[test_size_original, test_size_transfer]  # Lengths of the resulting splits\n",
-    ")\n",
+    "    # Split the training dataset into training and validation sets\n",
+    "    train_size = int(0.9 * len(train_val_dataset))  # Calculate the size of the training set (90% of the original)\n",
+    "    val_size = len(train_val_dataset) - train_size  # Calculate the size of the validation set (remaining 10%)\n",
+    "    train_dataset, val_dataset = torch.utils.data.random_split(\n",
+    "        dataset=train_val_dataset,  # Original training dataset to split\n",
+    "        lengths=[train_size, val_size]  # Lengths of the resulting splits\n",
+    "    )\n",
     "\n",
-    "# Display the training dataset object\n",
-    "train_dataset\n",
+    "    # Split the test dataset into two halves: original and transfer sets\n",
+    "    test_size_original = int(0.5 * len(test_dataset))  # Calculate the size of the original test set (50% of the original)\n",
+    "    test_size_transfer = len(test_dataset) - test_size_original  # Calculate the size of the transfer test set (remaining 50%)\n",
+    "    test_dataset_original, test_dataset_transfer = torch.utils.data.random_split(\n",
+    "        dataset=test_dataset,  # Original test dataset to split\n",
+    "        lengths=[test_size_original, test_size_transfer]  # Lengths of the resulting splits\n",
+    "    )\n",
     "\n",
-    "# Print the sizes of the training, validation, original test, and transfer test datasets\n",
-    "len(train_dataset), len(val_dataset), len(test_dataset_original), len(test_dataset_transfer)"
+    "    # Display the training dataset object\n",
+    "    train_dataset"
    ]
   },
   {
diff --git a/tutorials/W1D2_ComparingTasks/W1D2_Tutorial2.ipynb b/tutorials/W1D2_ComparingTasks/W1D2_Tutorial2.ipynb
index b38c35e13..77c87c98f 100644
--- a/tutorials/W1D2_ComparingTasks/W1D2_Tutorial2.ipynb
+++ b/tutorials/W1D2_ComparingTasks/W1D2_Tutorial2.ipynb
@@ -84,8 +84,7 @@
    "cell_type": "markdown",
    "id": "273ae64d-bc44-4e0e-b077-487da8391334",
    "metadata": {
-    "execution": {},
-    "jp-MarkdownHeadingCollapsed": true
+    "execution": {}
    },
    "source": [
     "---\n",
@@ -153,6 +152,8 @@
     "\n",
     "import logging\n",
     "import gc\n",
+    "import contextlib\n",
+    "import io\n",
     "\n",
     "# PyTorch and related libraries\n",
     "import torch\n",
@@ -181,8 +182,7 @@
    "cell_type": "markdown",
    "id": "910f242f",
    "metadata": {
-    "execution": {},
-    "jp-MarkdownHeadingCollapsed": true
+    "execution": {}
    },
    "source": [
     "## Figure settings\n"
@@ -211,40 +211,11 @@
     "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/course-content/main/nma.mplstyle\")"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "id": "01471a7f",
-   "metadata": {
-    "execution": {},
-    "jp-MarkdownHeadingCollapsed": true
-   },
-   "source": [
-    "##  Plotting functions\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "598cc45f-c4b7-406c-a80b-6adfae387583",
-   "metadata": {
-    "cellView": "form",
-    "execution": {},
-    "tags": [
-     "hide-input"
-    ]
-   },
-   "outputs": [],
-   "source": [
-    "# @title Plotting functions\n",
-    "# @markdown"
-   ]
-  },
   {
    "cell_type": "markdown",
    "id": "808d728a",
    "metadata": {
-    "execution": {},
-    "jp-MarkdownHeadingCollapsed": true
+    "execution": {}
    },
    "source": [
     "##  Helper functions"
@@ -526,8 +497,9 @@
     "    torchvision.transforms.Normalize((0.1307,), (0.3081,))  # Normalize the images with mean and standard deviation\n",
     "])\n",
     "\n",
-    "# Load the MNIST test dataset with the defined transformations\n",
-    "test_dset = torchvision.datasets.MNIST(\"./\", train=False, transform=mnist_transforms, download=True)\n",
+    "with contextlib.redirect_stdout(io.StringIO()):\n",
+    "    # Load the MNIST test dataset with the defined transformations\n",
+    "    test_dset = torchvision.datasets.MNIST(\"./\", train=False, transform=mnist_transforms, download=True)\n",
     "\n",
     "# Calculate the height and width of the MNIST images (28x28)\n",
     "height = int(784**0.5)\n",
diff --git a/tutorials/W1D2_ComparingTasks/instructor/W1D2_Tutorial1.ipynb b/tutorials/W1D2_ComparingTasks/instructor/W1D2_Tutorial1.ipynb
index 404f422db..c9d6fc0eb 100644
--- a/tutorials/W1D2_ComparingTasks/instructor/W1D2_Tutorial1.ipynb
+++ b/tutorials/W1D2_ComparingTasks/instructor/W1D2_Tutorial1.ipynb
@@ -176,6 +176,8 @@
     "from pathlib import Path\n",
     "import zipfile\n",
     "import random\n",
+    "import contextlib\n",
+    "import io\n",
     "\n",
     "# Import third-party libraries\n",
     "import numpy as np\n",
@@ -844,50 +846,49 @@
    },
    "outputs": [],
    "source": [
-    "# Define a transformation pipeline for the MNIST dataset\n",
-    "mnist_transform = transforms.Compose([\n",
-    "    transforms.Resize((32, 32)),  # Resize the images to 32x32 pixels\n",
-    "    transforms.ToTensor(),  # Convert images to PyTorch tensors\n",
-    "    transforms.Normalize(mean=(0.1307,), std=(0.3081,))  # Normalize the images with mean and standard deviation\n",
-    "])\n",
-    "\n",
-    "# Load the MNIST training dataset with transformations applied\n",
-    "train_val_dataset = torchvision.datasets.MNIST(\n",
-    "    root='./data',  # Directory to store/load the data\n",
-    "    train=True,  # Specify to load the training set\n",
-    "    transform=mnist_transform,  # Apply the transformation pipeline defined earlier\n",
-    "    download=True  # Download the dataset if it's not already present\n",
-    ")\n",
-    "\n",
-    "# Load the MNIST test dataset with transformations applied\n",
-    "test_dataset = torchvision.datasets.MNIST(\n",
-    "    root='./data',  # Directory to store/load the data\n",
-    "    train=False,  # Specify to load the test set\n",
-    "    transform=mnist_transform,  # Apply the transformation pipeline defined earlier\n",
-    "    download=True  # Download the dataset if it's not already present\n",
-    ")\n",
+    "with contextlib.redirect_stdout(io.StringIO()):\n",
+    "\n",
+    "    # Define a transformation pipeline for the MNIST dataset\n",
+    "    mnist_transform = transforms.Compose([\n",
+    "        transforms.Resize((32, 32)),  # Resize the images to 32x32 pixels\n",
+    "        transforms.ToTensor(),  # Convert images to PyTorch tensors\n",
+    "        transforms.Normalize(mean=(0.1307,), std=(0.3081,))  # Normalize the images with mean and standard deviation\n",
+    "    ])\n",
+    "\n",
+    "    # Load the MNIST training dataset with transformations applied\n",
+    "    train_val_dataset = torchvision.datasets.MNIST(\n",
+    "        root='./data',  # Directory to store/load the data\n",
+    "        train=True,  # Specify to load the training set\n",
+    "        transform=mnist_transform,  # Apply the transformation pipeline defined earlier\n",
+    "        download=True  # Download the dataset if it's not already present\n",
+    "    )\n",
     "\n",
-    "# Split the training dataset into training and validation sets\n",
-    "train_size = int(0.9 * len(train_val_dataset))  # Calculate the size of the training set (90% of the original)\n",
-    "val_size = len(train_val_dataset) - train_size  # Calculate the size of the validation set (remaining 10%)\n",
-    "train_dataset, val_dataset = torch.utils.data.random_split(\n",
-    "    dataset=train_val_dataset,  # Original training dataset to split\n",
-    "    lengths=[train_size, val_size]  # Lengths of the resulting splits\n",
-    ")\n",
+    "    # Load the MNIST test dataset with transformations applied\n",
+    "    test_dataset = torchvision.datasets.MNIST(\n",
+    "        root='./data',  # Directory to store/load the data\n",
+    "        train=False,  # Specify to load the test set\n",
+    "        transform=mnist_transform,  # Apply the transformation pipeline defined earlier\n",
+    "        download=True  # Download the dataset if it's not already present\n",
+    "    )\n",
     "\n",
-    "# Split the test dataset into two halves: original and transfer sets\n",
-    "test_size_original = int(0.5 * len(test_dataset))  # Calculate the size of the original test set (50% of the original)\n",
-    "test_size_transfer = len(test_dataset) - test_size_original  # Calculate the size of the transfer test set (remaining 50%)\n",
-    "test_dataset_original, test_dataset_transfer = torch.utils.data.random_split(\n",
-    "    dataset=test_dataset,  # Original test dataset to split\n",
-    "    lengths=[test_size_original, test_size_transfer]  # Lengths of the resulting splits\n",
-    ")\n",
+    "    # Split the training dataset into training and validation sets\n",
+    "    train_size = int(0.9 * len(train_val_dataset))  # Calculate the size of the training set (90% of the original)\n",
+    "    val_size = len(train_val_dataset) - train_size  # Calculate the size of the validation set (remaining 10%)\n",
+    "    train_dataset, val_dataset = torch.utils.data.random_split(\n",
+    "        dataset=train_val_dataset,  # Original training dataset to split\n",
+    "        lengths=[train_size, val_size]  # Lengths of the resulting splits\n",
+    "    )\n",
     "\n",
-    "# Display the training dataset object\n",
-    "train_dataset\n",
+    "    # Split the test dataset into two halves: original and transfer sets\n",
+    "    test_size_original = int(0.5 * len(test_dataset))  # Calculate the size of the original test set (50% of the original)\n",
+    "    test_size_transfer = len(test_dataset) - test_size_original  # Calculate the size of the transfer test set (remaining 50%)\n",
+    "    test_dataset_original, test_dataset_transfer = torch.utils.data.random_split(\n",
+    "        dataset=test_dataset,  # Original test dataset to split\n",
+    "        lengths=[test_size_original, test_size_transfer]  # Lengths of the resulting splits\n",
+    "    )\n",
     "\n",
-    "# Print the sizes of the training, validation, original test, and transfer test datasets\n",
-    "len(train_dataset), len(val_dataset), len(test_dataset_original), len(test_dataset_transfer)"
+    "    # Display the training dataset object\n",
+    "    train_dataset"
    ]
   },
   {
diff --git a/tutorials/W1D2_ComparingTasks/instructor/W1D2_Tutorial2.ipynb b/tutorials/W1D2_ComparingTasks/instructor/W1D2_Tutorial2.ipynb
index 7340891b8..630159fba 100644
--- a/tutorials/W1D2_ComparingTasks/instructor/W1D2_Tutorial2.ipynb
+++ b/tutorials/W1D2_ComparingTasks/instructor/W1D2_Tutorial2.ipynb
@@ -84,8 +84,7 @@
    "cell_type": "markdown",
    "id": "273ae64d-bc44-4e0e-b077-487da8391334",
    "metadata": {
-    "execution": {},
-    "jp-MarkdownHeadingCollapsed": true
+    "execution": {}
    },
    "source": [
     "---\n",
@@ -153,6 +152,8 @@
     "\n",
     "import logging\n",
     "import gc\n",
+    "import contextlib\n",
+    "import io\n",
     "\n",
     "# PyTorch and related libraries\n",
     "import torch\n",
@@ -181,8 +182,7 @@
    "cell_type": "markdown",
    "id": "910f242f",
    "metadata": {
-    "execution": {},
-    "jp-MarkdownHeadingCollapsed": true
+    "execution": {}
    },
    "source": [
     "## Figure settings\n"
@@ -211,40 +211,11 @@
     "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/course-content/main/nma.mplstyle\")"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "id": "01471a7f",
-   "metadata": {
-    "execution": {},
-    "jp-MarkdownHeadingCollapsed": true
-   },
-   "source": [
-    "##  Plotting functions\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "598cc45f-c4b7-406c-a80b-6adfae387583",
-   "metadata": {
-    "cellView": "form",
-    "execution": {},
-    "tags": [
-     "hide-input"
-    ]
-   },
-   "outputs": [],
-   "source": [
-    "# @title Plotting functions\n",
-    "# @markdown"
-   ]
-  },
   {
    "cell_type": "markdown",
    "id": "808d728a",
    "metadata": {
-    "execution": {},
-    "jp-MarkdownHeadingCollapsed": true
+    "execution": {}
    },
    "source": [
     "##  Helper functions"
@@ -526,8 +497,9 @@
     "    torchvision.transforms.Normalize((0.1307,), (0.3081,))  # Normalize the images with mean and standard deviation\n",
     "])\n",
     "\n",
-    "# Load the MNIST test dataset with the defined transformations\n",
-    "test_dset = torchvision.datasets.MNIST(\"./\", train=False, transform=mnist_transforms, download=True)\n",
+    "with contextlib.redirect_stdout(io.StringIO()):\n",
+    "    # Load the MNIST test dataset with the defined transformations\n",
+    "    test_dset = torchvision.datasets.MNIST(\"./\", train=False, transform=mnist_transforms, download=True)\n",
     "\n",
     "# Calculate the height and width of the MNIST images (28x28)\n",
     "height = int(784**0.5)\n",
diff --git a/tutorials/W1D2_ComparingTasks/static/W1D2_Tutorial2_Solution_39154423_0.png b/tutorials/W1D2_ComparingTasks/static/W1D2_Tutorial2_Solution_39154423_0.png
index c1587f598..bb397ef85 100644
Binary files a/tutorials/W1D2_ComparingTasks/static/W1D2_Tutorial2_Solution_39154423_0.png and b/tutorials/W1D2_ComparingTasks/static/W1D2_Tutorial2_Solution_39154423_0.png differ
diff --git a/tutorials/W1D2_ComparingTasks/student/W1D2_Tutorial1.ipynb b/tutorials/W1D2_ComparingTasks/student/W1D2_Tutorial1.ipynb
index 5a2ecc3f8..d54c0a9a7 100644
--- a/tutorials/W1D2_ComparingTasks/student/W1D2_Tutorial1.ipynb
+++ b/tutorials/W1D2_ComparingTasks/student/W1D2_Tutorial1.ipynb
@@ -176,6 +176,8 @@
     "from pathlib import Path\n",
     "import zipfile\n",
     "import random\n",
+    "import contextlib\n",
+    "import io\n",
     "\n",
     "# Import third-party libraries\n",
     "import numpy as np\n",
@@ -844,50 +846,49 @@
    },
    "outputs": [],
    "source": [
-    "# Define a transformation pipeline for the MNIST dataset\n",
-    "mnist_transform = transforms.Compose([\n",
-    "    transforms.Resize((32, 32)),  # Resize the images to 32x32 pixels\n",
-    "    transforms.ToTensor(),  # Convert images to PyTorch tensors\n",
-    "    transforms.Normalize(mean=(0.1307,), std=(0.3081,))  # Normalize the images with mean and standard deviation\n",
-    "])\n",
-    "\n",
-    "# Load the MNIST training dataset with transformations applied\n",
-    "train_val_dataset = torchvision.datasets.MNIST(\n",
-    "    root='./data',  # Directory to store/load the data\n",
-    "    train=True,  # Specify to load the training set\n",
-    "    transform=mnist_transform,  # Apply the transformation pipeline defined earlier\n",
-    "    download=True  # Download the dataset if it's not already present\n",
-    ")\n",
-    "\n",
-    "# Load the MNIST test dataset with transformations applied\n",
-    "test_dataset = torchvision.datasets.MNIST(\n",
-    "    root='./data',  # Directory to store/load the data\n",
-    "    train=False,  # Specify to load the test set\n",
-    "    transform=mnist_transform,  # Apply the transformation pipeline defined earlier\n",
-    "    download=True  # Download the dataset if it's not already present\n",
-    ")\n",
+    "with contextlib.redirect_stdout(io.StringIO()):\n",
+    "\n",
+    "    # Define a transformation pipeline for the MNIST dataset\n",
+    "    mnist_transform = transforms.Compose([\n",
+    "        transforms.Resize((32, 32)),  # Resize the images to 32x32 pixels\n",
+    "        transforms.ToTensor(),  # Convert images to PyTorch tensors\n",
+    "        transforms.Normalize(mean=(0.1307,), std=(0.3081,))  # Normalize the images with mean and standard deviation\n",
+    "    ])\n",
+    "\n",
+    "    # Load the MNIST training dataset with transformations applied\n",
+    "    train_val_dataset = torchvision.datasets.MNIST(\n",
+    "        root='./data',  # Directory to store/load the data\n",
+    "        train=True,  # Specify to load the training set\n",
+    "        transform=mnist_transform,  # Apply the transformation pipeline defined earlier\n",
+    "        download=True  # Download the dataset if it's not already present\n",
+    "    )\n",
     "\n",
-    "# Split the training dataset into training and validation sets\n",
-    "train_size = int(0.9 * len(train_val_dataset))  # Calculate the size of the training set (90% of the original)\n",
-    "val_size = len(train_val_dataset) - train_size  # Calculate the size of the validation set (remaining 10%)\n",
-    "train_dataset, val_dataset = torch.utils.data.random_split(\n",
-    "    dataset=train_val_dataset,  # Original training dataset to split\n",
-    "    lengths=[train_size, val_size]  # Lengths of the resulting splits\n",
-    ")\n",
+    "    # Load the MNIST test dataset with transformations applied\n",
+    "    test_dataset = torchvision.datasets.MNIST(\n",
+    "        root='./data',  # Directory to store/load the data\n",
+    "        train=False,  # Specify to load the test set\n",
+    "        transform=mnist_transform,  # Apply the transformation pipeline defined earlier\n",
+    "        download=True  # Download the dataset if it's not already present\n",
+    "    )\n",
     "\n",
-    "# Split the test dataset into two halves: original and transfer sets\n",
-    "test_size_original = int(0.5 * len(test_dataset))  # Calculate the size of the original test set (50% of the original)\n",
-    "test_size_transfer = len(test_dataset) - test_size_original  # Calculate the size of the transfer test set (remaining 50%)\n",
-    "test_dataset_original, test_dataset_transfer = torch.utils.data.random_split(\n",
-    "    dataset=test_dataset,  # Original test dataset to split\n",
-    "    lengths=[test_size_original, test_size_transfer]  # Lengths of the resulting splits\n",
-    ")\n",
+    "    # Split the training dataset into training and validation sets\n",
+    "    train_size = int(0.9 * len(train_val_dataset))  # Calculate the size of the training set (90% of the original)\n",
+    "    val_size = len(train_val_dataset) - train_size  # Calculate the size of the validation set (remaining 10%)\n",
+    "    train_dataset, val_dataset = torch.utils.data.random_split(\n",
+    "        dataset=train_val_dataset,  # Original training dataset to split\n",
+    "        lengths=[train_size, val_size]  # Lengths of the resulting splits\n",
+    "    )\n",
     "\n",
-    "# Display the training dataset object\n",
-    "train_dataset\n",
+    "    # Split the test dataset into two halves: original and transfer sets\n",
+    "    test_size_original = int(0.5 * len(test_dataset))  # Calculate the size of the original test set (50% of the original)\n",
+    "    test_size_transfer = len(test_dataset) - test_size_original  # Calculate the size of the transfer test set (remaining 50%)\n",
+    "    test_dataset_original, test_dataset_transfer = torch.utils.data.random_split(\n",
+    "        dataset=test_dataset,  # Original test dataset to split\n",
+    "        lengths=[test_size_original, test_size_transfer]  # Lengths of the resulting splits\n",
+    "    )\n",
     "\n",
-    "# Print the sizes of the training, validation, original test, and transfer test datasets\n",
-    "len(train_dataset), len(val_dataset), len(test_dataset_original), len(test_dataset_transfer)"
+    "    # Display the training dataset object\n",
+    "    train_dataset"
    ]
   },
   {
diff --git a/tutorials/W1D2_ComparingTasks/student/W1D2_Tutorial2.ipynb b/tutorials/W1D2_ComparingTasks/student/W1D2_Tutorial2.ipynb
index 2da3b1acf..9316745d3 100644
--- a/tutorials/W1D2_ComparingTasks/student/W1D2_Tutorial2.ipynb
+++ b/tutorials/W1D2_ComparingTasks/student/W1D2_Tutorial2.ipynb
@@ -84,8 +84,7 @@
    "cell_type": "markdown",
    "id": "273ae64d-bc44-4e0e-b077-487da8391334",
    "metadata": {
-    "execution": {},
-    "jp-MarkdownHeadingCollapsed": true
+    "execution": {}
    },
    "source": [
     "---\n",
@@ -153,6 +152,8 @@
     "\n",
     "import logging\n",
     "import gc\n",
+    "import contextlib\n",
+    "import io\n",
     "\n",
     "# PyTorch and related libraries\n",
     "import torch\n",
@@ -181,8 +182,7 @@
    "cell_type": "markdown",
    "id": "910f242f",
    "metadata": {
-    "execution": {},
-    "jp-MarkdownHeadingCollapsed": true
+    "execution": {}
    },
    "source": [
     "## Figure settings\n"
@@ -211,40 +211,11 @@
     "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/course-content/main/nma.mplstyle\")"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "id": "01471a7f",
-   "metadata": {
-    "execution": {},
-    "jp-MarkdownHeadingCollapsed": true
-   },
-   "source": [
-    "##  Plotting functions\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "598cc45f-c4b7-406c-a80b-6adfae387583",
-   "metadata": {
-    "cellView": "form",
-    "execution": {},
-    "tags": [
-     "hide-input"
-    ]
-   },
-   "outputs": [],
-   "source": [
-    "# @title Plotting functions\n",
-    "# @markdown"
-   ]
-  },
   {
    "cell_type": "markdown",
    "id": "808d728a",
    "metadata": {
-    "execution": {},
-    "jp-MarkdownHeadingCollapsed": true
+    "execution": {}
    },
    "source": [
     "##  Helper functions"
@@ -526,8 +497,9 @@
     "    torchvision.transforms.Normalize((0.1307,), (0.3081,))  # Normalize the images with mean and standard deviation\n",
     "])\n",
     "\n",
-    "# Load the MNIST test dataset with the defined transformations\n",
-    "test_dset = torchvision.datasets.MNIST(\"./\", train=False, transform=mnist_transforms, download=True)\n",
+    "with contextlib.redirect_stdout(io.StringIO()):\n",
+    "    # Load the MNIST test dataset with the defined transformations\n",
+    "    test_dset = torchvision.datasets.MNIST(\"./\", train=False, transform=mnist_transforms, download=True)\n",
     "\n",
     "# Calculate the height and width of the MNIST images (28x28)\n",
     "height = int(784**0.5)\n",