diff --git a/exercise.ipynb b/exercise.ipynb
index 0ad1bfb..bde574d 100644
--- a/exercise.ipynb
+++ b/exercise.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "0e557335",
+   "id": "eeffd47d",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -29,7 +29,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "8667a26e",
+   "id": "d9b65175",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -41,7 +41,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e36eca8a",
+   "id": "025f7566",
    "metadata": {},
    "source": [
     "\n",
@@ -54,7 +54,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ffea56f9",
+   "id": "7f4771f3",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -68,7 +68,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "55cab99a",
+   "id": "9d60bf35",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -84,7 +84,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1fe21cf9",
+   "id": "4acd91e9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -102,35 +102,35 @@
   },
   {
    "cell_type": "markdown",
-   "id": "610b5cdb",
+   "id": "9e70adfb",
    "metadata": {
     "lines_to_next_cell": 0
    },
    "source": [
-    "We have pre-traiend a classifier for you on this dataset. It is the same architecture classifier as you used in the Failure Modes exercise: a `DenseModel`.\n",
+    "During the setup you have pre-traiend a classifier on this dataset. It is the same architecture classifier as you used in the Failure Modes exercise: a `DenseModel`.\n",
     "Let's load that classifier now!"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "275b9fc0",
+   "id": "aa1a58cf",
    "metadata": {
     "lines_to_next_cell": 0
    },
    "source": [
     "<div class=\"alert alert-block alert-info\"><h3>Task 1.1: Load the classifier</h3>\n",
-    "We have written a slightly more general version of the `DenseModel` that you used in the previous exercise. Ours requires two inputs:\n",
-    "- `input_shape`: the shape of the input images, as a tuple\n",
-    "- `num_classes`: the number of classes in the dataset\n",
+    "We have written a slightly more general version of the <code>DenseModel</code> that you used in the previous exercise. Ours requires two inputs:\n",
+    "<li> <code>input_shape</code>: the shape of the input images, as a tuple </li>\n",
+    "<li> <code>num_classes</code>: the number of classes in the dataset </li>\n",
     "\n",
     "Create a dense model with the right inputs and load the weights from the checkpoint.\n",
-    "<div>"
+    "</div>"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a6ce29c6",
+   "id": "fb0683a8",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": [
@@ -155,7 +155,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3f2739bc",
+   "id": "a3326bfe",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -166,7 +166,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9423e54b",
+   "id": "19354115",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -193,7 +193,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "cecb6c11",
+   "id": "ff293353",
    "metadata": {},
    "source": [
     "# Part 2: Using Integrated Gradients to find what the classifier knows\n",
@@ -203,7 +203,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c57a8a9d",
+   "id": "c8f3e86d",
    "metadata": {},
    "source": [
     "## Attributions through integrated gradients\n",
@@ -216,7 +216,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a596affd",
+   "id": "c6071dca",
    "metadata": {
     "tags": []
    },
@@ -234,7 +234,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "06db8b81",
+   "id": "7be4fc86",
    "metadata": {
     "tags": []
    },
@@ -250,7 +250,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a8252672",
+   "id": "4766f414",
    "metadata": {
     "tags": [
      "task"
@@ -271,7 +271,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "115b94e7",
+   "id": "5395d4a2",
    "metadata": {
     "tags": []
    },
@@ -284,7 +284,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6e281067",
+   "id": "f75d2492",
    "metadata": {
     "lines_to_next_cell": 2,
     "tags": []
@@ -296,7 +296,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "07528ee6",
+   "id": "9d788363",
    "metadata": {
     "tags": []
    },
@@ -324,7 +324,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "780d14de",
+   "id": "56d5e9f3",
    "metadata": {
     "tags": []
    },
@@ -337,13 +337,13 @@
   },
   {
    "cell_type": "markdown",
-   "id": "33d9db66",
+   "id": "4c28aff0",
    "metadata": {
     "lines_to_next_cell": 2
    },
    "source": [
     "\n",
-    "The attributions are shown as a heatmap. The brighter the pixel, the more important this attribution method thinks that it is.\n",
+    "The attributions are shown as a heatmap. The closer to 1 the pixel value, the more important this attribution method thinks that it is.\n",
     "As you can see, it is pretty good at recognizing the number within the image.\n",
     "As we know, however, it is not the digit itself that is important for the classification, it is the color!\n",
     "Although the method is picking up really well on the region of interest, it would be difficult to conclude from this that it is the color that matters."
@@ -351,7 +351,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "81725a7f",
+   "id": "e871c574",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -364,7 +364,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "cba802d4",
+   "id": "9bd13121",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -389,13 +389,13 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7254eab3",
+   "id": "e302a5ca",
    "metadata": {
     "lines_to_next_cell": 0
    },
    "source": [
     "We get some better clues when looking at the attributions in color.\n",
-    "The highlighting doesn't just happen in the region with number, but also seems to hapen in a channel that matches the color of the image.\n",
+    "The highlighting doesn't just happen in the region with number, but also seems to happen in a channel that matches the color of the image.\n",
     "Just based on this, however, we don't get much more information than we got from the images themselves.\n",
     "\n",
     "If we didn't know in advance, it is unclear whether the color or the number is the most important feature for the classifier."
@@ -403,7 +403,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2fd14add",
+   "id": "844f37f0",
    "metadata": {},
    "source": [
     "\n",
@@ -429,7 +429,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "9e88a188",
+   "id": "d6e8e759",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-info\"><h4>Task 2.3: Use random noise as a baseline</h4>\n",
@@ -441,7 +441,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "94ae54d4",
+   "id": "74344ee8",
    "metadata": {
     "tags": [
      "task"
@@ -453,16 +453,17 @@
     "random_baselines = ...  # TODO Change\n",
     "# Generate the attributions\n",
     "attributions_random = integrated_gradients.attribute(...)  # TODO Change\n",
+    "attributions_random = attributions_random.cpu().numpy()\n",
     "\n",
     "# Plotting\n",
-    "for attr, im, lbl in zip(attributions, x.cpu().numpy(), y.cpu().numpy()):\n",
+    "for attr, im, lbl in zip(attributions_random, x.cpu().numpy(), y.cpu().numpy()):\n",
     "    print(f\"Class {lbl}\")\n",
-    "    visualize_attribution(attr, im)"
+    "    visualize_color_attribution(attr, im)"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "8d7fbe0e",
+   "id": "6d560996",
    "metadata": {
     "tags": []
    },
@@ -476,7 +477,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "438f77b4",
+   "id": "e15c02f7",
    "metadata": {
     "tags": [
      "task"
@@ -491,15 +492,16 @@
     "# Generate the attributions\n",
     "attributions_blurred = integrated_gradients.attribute(...)  # TODO Fill\n",
     "\n",
+    "attributions_blurred = attributions_blurred.cpu().numpy()\n",
     "# Plotting\n",
-    "for attr, im, lbl in zip(attributions, x.cpu().numpy(), y.cpu().numpy()):\n",
+    "for attr, im, lbl in zip(attributions_blurred, x.cpu().numpy(), y.cpu().numpy()):\n",
     "    print(f\"Class {lbl}\")\n",
     "    visualize_color_attribution(attr, im)"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "220cc456",
+   "id": "de810272",
    "metadata": {
     "tags": []
    },
@@ -515,7 +517,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "83bbf9b4",
+   "id": "7a6f89dc",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-info\"><h2>BONUS Task: Using different attributions.</h2>\n",
@@ -529,11 +531,11 @@
   },
   {
    "cell_type": "markdown",
-   "id": "69d66239",
+   "id": "f3e293c5",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-success\"><h2>Checkpoint 2</h2>\n",
-    "Let us know on the exercise chat when you've reached this point!\n",
+    "Put up your green sticky note when you've reached this point!\n",
     "\n",
     "At this point we have:\n",
     "\n",
@@ -549,7 +551,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0d20d9cf",
+   "id": "fb3048e1",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -562,9 +564,9 @@
     "**What is a counterfactual?**\n",
     "\n",
     "You've learned about adversarial examples in the lecture on failure modes. These are the imperceptible or noisy changes to an image that drastically changes a classifier's opinion.\n",
-    "Counterfactual explanations are the useful cousins of adversarial examples. They are *perceptible* and *informative* changes to an image that changes a classifier's opinion.\n",
+    "Counterfactual explanations are the useful cousins of adversarial examples. They are *perceptible* and *informative* changes to an image that change a classifier's opinion.\n",
     "\n",
-    "In the image below you can see the difference between the two. In the first column are MNIST images along with their classifictaions, and in the second column are counterfactual explanations to *change* that class. You can see that in both cases a human being would (hopefully) agree with the new classification. By comparing the two columns, we can therefore begin to define what makes each digit special.\n",
+    "In the image below you can see the difference between the two. In the first column are (non-color) MNIST images along with their classifications, and in the second column are counterfactual explanations to *change* that class. You can see that in both cases a human being would (hopefully) agree with the new classification. By comparing the two columns, we can therefore begin to define what makes each digit special.\n",
     "\n",
     "In contrast, the third and fourth columns show an MNIST image and a corresponding adversarial example. Here the network returns a prediction that most human beings (who aren't being facetious) would strongly disagree with.\n",
     "\n",
@@ -577,7 +579,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "18cda8e3",
+   "id": "88f77592",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -600,7 +602,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a4ea7f03",
+   "id": "6742eca7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -632,7 +634,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "53e15fe2",
+   "id": "3ed703e3",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -647,7 +649,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "35e552da",
+   "id": "57a7d702",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": [
@@ -656,7 +658,7 @@
    },
    "outputs": [],
    "source": [
-    "style_size = ...  # TODO choose a size for the style space\n",
+    "style_size = 3\n",
     "unet_depth = ...  # TODO Choose a depth for the UNet\n",
     "style_encoder = DenseModel(\n",
     "    input_shape=..., num_classes=...  # How big is the style space?\n",
@@ -668,14 +670,14 @@
   },
   {
    "cell_type": "markdown",
-   "id": "49a9a89d",
+   "id": "fdf97172",
    "metadata": {
     "tags": []
    },
    "source": [
     "<div class=\"alert alert-block alert-warning\"><h3>Hyper-parameter choices</h3>\n",
     "<ul>\n",
-    "<li>Are any of the hyperparameters you choose above constrained in some way?</li>\n",
+    "<li>Are any of the hyperparameters above constrained in some way?</li>\n",
     "<li>What would happen if you chose a depth of 10 for the UNet?</li>\n",
     "<li>Is there a minimum size for the style space? Why or why not?</li>\n",
     "</ul>"
@@ -683,7 +685,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "fb9b4a59",
+   "id": "bcb08a78",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -700,7 +702,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "811aec8a",
+   "id": "349544f2",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": [
@@ -714,7 +716,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ef8dd1e5",
+   "id": "fd541400",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -725,7 +727,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "069f2804",
+   "id": "1e91e416",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -735,7 +737,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "380256c9",
+   "id": "b0377825",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -753,7 +755,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "576200cc",
+   "id": "f4939a8d",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -765,7 +767,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2885a3a0",
+   "id": "9a6ca642",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -784,7 +786,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9cc5eb23",
+   "id": "382fa57e",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -793,7 +795,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5eb04071",
+   "id": "ccf3b275",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -809,7 +811,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "805dfbf2",
+   "id": "0a424199",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -818,7 +820,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "8fab1a3e",
+   "id": "fe335f99",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -830,7 +832,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "dc899114",
+   "id": "977ce4c8",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -843,7 +845,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7e13436e",
+   "id": "7bf8cb24",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -857,7 +859,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1cb60fb1",
+   "id": "0d5ee1e1",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -869,7 +871,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e8441eda",
+   "id": "43c93243",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -889,7 +891,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6bb72378",
+   "id": "9a789ba4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -913,7 +915,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "af5d018a",
+   "id": "91b15866",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -923,7 +925,26 @@
   },
   {
    "cell_type": "markdown",
-   "id": "97aaeef6",
+   "id": "87d5a8a1",
+   "metadata": {},
+   "source": [
+    "<div class=\"alert alert-block alert-success\"><h2>Checkpoint 3</h2>\n",
+    "Put up your green sticky note when you've reached this point!\n",
+    "\n",
+    "At this point we have:\n",
+    "\n",
+    "- Loaded a classifier that classifies MNIST-like images by color, but we don't know how!\n",
+    "- Tried applying Integrated Gradients to find out what the classifier is looking at - with little success.\n",
+    "- Discovered the effect of changing the baseline on the output of integrated gradients.\n",
+    "- Defined the hyperparameters for a StarGAN to create counterfactual images.\n",
+    "\n",
+    "Next up, we will define the training loop for the StarGAN.\n",
+    "</div>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3a1e7dd9",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -945,7 +966,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "74570a78",
+   "id": "c6348be7",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -957,7 +978,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2b7a8a55",
+   "id": "f5ed4de5",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": [
@@ -1068,7 +1089,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "41e0e738",
+   "id": "f6cd135b",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1080,7 +1101,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b49fd330",
+   "id": "cdecbbbe",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1096,7 +1117,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "77d4f778",
+   "id": "2c7d13e4",
    "metadata": {
     "tags": []
    },
@@ -1111,7 +1132,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2258df98",
+   "id": "689db81d",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1123,7 +1144,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "fa248f8b",
+   "id": "c3f514fa",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1145,7 +1166,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c2db79f3",
+   "id": "27ee28e5",
    "metadata": {
     "tags": []
    },
@@ -1155,13 +1176,13 @@
     "The same method can be used to create a StarGAN with different basic elements.\n",
     "For example, you can change the archictecture of the generators, or of the discriminator to better fit your data in the future.\n",
     "\n",
-    "You know the drill... let us know on the exercise chat when you have arrived here!\n",
+    "You know the drill... put up your green sticky note when you have arrived here!\n",
     "</div>"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "de5cb982",
+   "id": "e1c592af",
    "metadata": {
     "tags": []
    },
@@ -1171,7 +1192,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ac566b2e",
+   "id": "d2374be1",
    "metadata": {
     "tags": []
    },
@@ -1188,7 +1209,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "af5bbb66",
+   "id": "c038ee90",
    "metadata": {
     "title": "Loading the test dataset"
    },
@@ -1208,7 +1229,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e5392251",
+   "id": "2dd7e283",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1220,7 +1241,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "02808bdf",
+   "id": "1910ed0f",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1233,7 +1254,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e2e7f289",
+   "id": "99188e3e",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1243,7 +1264,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "9b936567",
+   "id": "fe2b4929",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1261,7 +1282,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7e5d760d",
+   "id": "30cd9209",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": [
@@ -1297,7 +1318,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "35dd9913",
+   "id": "ef56e941",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1309,7 +1330,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "13fa9bfa",
+   "id": "1612f8ce",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1322,7 +1343,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "449428b9",
+   "id": "2239beb9",
    "metadata": {
     "tags": []
    },
@@ -1337,7 +1358,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "cb3d6985",
+   "id": "168352ff",
    "metadata": {
     "tags": []
    },
@@ -1348,7 +1369,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3096db55",
+   "id": "f24cdc48",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1362,7 +1383,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "19fd90a6",
+   "id": "b093ffc8",
    "metadata": {
     "tags": []
    },
@@ -1377,7 +1398,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e0967ec3",
+   "id": "be5c5a98",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1392,7 +1413,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b3312425",
+   "id": "0891fcf7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1413,7 +1434,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "42c617e2",
+   "id": "b080dfe1",
    "metadata": {
     "title": "Another visualization function"
    },
@@ -1442,7 +1463,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3075be99",
+   "id": "e41e00c9",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1458,7 +1479,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "d70f3ff4",
+   "id": "ac7678a7",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1474,7 +1495,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "593ba2db",
+   "id": "3e06986c",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1489,7 +1510,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "863c8dca",
+   "id": "7c9a99c4",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1503,7 +1524,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "602db341",
+   "id": "b1affa61",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1523,7 +1544,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "cf1e78db",
+   "id": "4bf0e982",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1559,7 +1580,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6e37fd94",
+   "id": "7f8f6ca5",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1573,7 +1594,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5b55e69e",
+   "id": "537edaa0",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1589,7 +1610,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3862ab46",
+   "id": "c4dd727d",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1605,7 +1626,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "bedeaf5c",
+   "id": "c5a5e2c0",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1628,7 +1649,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "75e28b2c",
+   "id": "a9ede52b",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-info\"><h3>Task 5.1: Explore the style space</h3>\n",
@@ -1640,7 +1661,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d2ffe8e5",
+   "id": "af4a9d95",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1675,7 +1696,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "559a66a2",
+   "id": "be742cb1",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1691,7 +1712,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4c268ad8",
+   "id": "09032a46",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1718,7 +1739,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7e9d9184",
+   "id": "4d9a7277",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1732,7 +1753,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "45005d21",
+   "id": "3016bf5c",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1749,7 +1770,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b6e19ec8",
+   "id": "b47bc9a2",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1772,7 +1793,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7aba6e43",
+   "id": "5434a70a",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1781,7 +1802,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c90115e2",
+   "id": "9d066df4",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-warning\"><h3>Questions</h3>\n",
@@ -1793,7 +1814,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6dd82862",
+   "id": "e4295029",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-success\"><h2>Checkpoint 5</h2>\n",
@@ -1811,7 +1832,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b6389689",
+   "id": "3cd7976f",
    "metadata": {
     "lines_to_next_cell": 0
    },
diff --git a/solution.ipynb b/solution.ipynb
index d22fd9a..9b17243 100644
--- a/solution.ipynb
+++ b/solution.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "0e557335",
+   "id": "eeffd47d",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -29,7 +29,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "8667a26e",
+   "id": "d9b65175",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -41,7 +41,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e36eca8a",
+   "id": "025f7566",
    "metadata": {},
    "source": [
     "\n",
@@ -54,7 +54,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ffea56f9",
+   "id": "7f4771f3",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -68,7 +68,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "55cab99a",
+   "id": "9d60bf35",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -84,7 +84,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1fe21cf9",
+   "id": "4acd91e9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -102,35 +102,35 @@
   },
   {
    "cell_type": "markdown",
-   "id": "610b5cdb",
+   "id": "9e70adfb",
    "metadata": {
     "lines_to_next_cell": 0
    },
    "source": [
-    "We have pre-traiend a classifier for you on this dataset. It is the same architecture classifier as you used in the Failure Modes exercise: a `DenseModel`.\n",
+    "During the setup you have pre-traiend a classifier on this dataset. It is the same architecture classifier as you used in the Failure Modes exercise: a `DenseModel`.\n",
     "Let's load that classifier now!"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "275b9fc0",
+   "id": "aa1a58cf",
    "metadata": {
     "lines_to_next_cell": 0
    },
    "source": [
     "<div class=\"alert alert-block alert-info\"><h3>Task 1.1: Load the classifier</h3>\n",
-    "We have written a slightly more general version of the `DenseModel` that you used in the previous exercise. Ours requires two inputs:\n",
-    "- `input_shape`: the shape of the input images, as a tuple\n",
-    "- `num_classes`: the number of classes in the dataset\n",
+    "We have written a slightly more general version of the <code>DenseModel</code> that you used in the previous exercise. Ours requires two inputs:\n",
+    "<li> <code>input_shape</code>: the shape of the input images, as a tuple </li>\n",
+    "<li> <code>num_classes</code>: the number of classes in the dataset </li>\n",
     "\n",
     "Create a dense model with the right inputs and load the weights from the checkpoint.\n",
-    "<div>"
+    "</div>"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "493d2b7b",
+   "id": "95d9fa51",
    "metadata": {
     "tags": [
      "solution"
@@ -154,7 +154,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3f2739bc",
+   "id": "a3326bfe",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -165,7 +165,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9423e54b",
+   "id": "19354115",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -192,7 +192,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "cecb6c11",
+   "id": "ff293353",
    "metadata": {},
    "source": [
     "# Part 2: Using Integrated Gradients to find what the classifier knows\n",
@@ -202,7 +202,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c57a8a9d",
+   "id": "c8f3e86d",
    "metadata": {},
    "source": [
     "## Attributions through integrated gradients\n",
@@ -215,7 +215,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a596affd",
+   "id": "c6071dca",
    "metadata": {
     "tags": []
    },
@@ -233,7 +233,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "06db8b81",
+   "id": "7be4fc86",
    "metadata": {
     "tags": []
    },
@@ -249,7 +249,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "78463dad",
+   "id": "e5f2cecc",
    "metadata": {
     "tags": [
      "solution"
@@ -273,7 +273,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "115b94e7",
+   "id": "5395d4a2",
    "metadata": {
     "tags": []
    },
@@ -286,7 +286,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6e281067",
+   "id": "f75d2492",
    "metadata": {
     "lines_to_next_cell": 2,
     "tags": []
@@ -298,7 +298,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "07528ee6",
+   "id": "9d788363",
    "metadata": {
     "tags": []
    },
@@ -326,7 +326,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "780d14de",
+   "id": "56d5e9f3",
    "metadata": {
     "tags": []
    },
@@ -339,13 +339,13 @@
   },
   {
    "cell_type": "markdown",
-   "id": "33d9db66",
+   "id": "4c28aff0",
    "metadata": {
     "lines_to_next_cell": 2
    },
    "source": [
     "\n",
-    "The attributions are shown as a heatmap. The brighter the pixel, the more important this attribution method thinks that it is.\n",
+    "The attributions are shown as a heatmap. The closer to 1 the pixel value, the more important this attribution method thinks that it is.\n",
     "As you can see, it is pretty good at recognizing the number within the image.\n",
     "As we know, however, it is not the digit itself that is important for the classification, it is the color!\n",
     "Although the method is picking up really well on the region of interest, it would be difficult to conclude from this that it is the color that matters."
@@ -353,7 +353,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "81725a7f",
+   "id": "e871c574",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -366,7 +366,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "cba802d4",
+   "id": "9bd13121",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -391,13 +391,13 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7254eab3",
+   "id": "e302a5ca",
    "metadata": {
     "lines_to_next_cell": 0
    },
    "source": [
     "We get some better clues when looking at the attributions in color.\n",
-    "The highlighting doesn't just happen in the region with number, but also seems to hapen in a channel that matches the color of the image.\n",
+    "The highlighting doesn't just happen in the region with number, but also seems to happen in a channel that matches the color of the image.\n",
     "Just based on this, however, we don't get much more information than we got from the images themselves.\n",
     "\n",
     "If we didn't know in advance, it is unclear whether the color or the number is the most important feature for the classifier."
@@ -405,7 +405,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2fd14add",
+   "id": "844f37f0",
    "metadata": {},
    "source": [
     "\n",
@@ -431,7 +431,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "9e88a188",
+   "id": "d6e8e759",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-info\"><h4>Task 2.3: Use random noise as a baseline</h4>\n",
@@ -443,7 +443,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3823dfba",
+   "id": "7135e83d",
    "metadata": {
     "tags": [
      "solution"
@@ -460,16 +460,16 @@
     "attributions_random = integrated_gradients.attribute(\n",
     "    x, target=y, baselines=random_baselines\n",
     ")\n",
-    "\n",
+    "attributions_random = attributions_random.cpu().numpy()\n",
     "# Plotting\n",
-    "for attr, im, lbl in zip(attributions, x.cpu().numpy(), y.cpu().numpy()):\n",
+    "for attr, im, lbl in zip(attributions_random, x.cpu().numpy(), y.cpu().numpy()):\n",
     "    print(f\"Class {lbl}\")\n",
     "    visualize_color_attribution(attr, im)"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "8d7fbe0e",
+   "id": "6d560996",
    "metadata": {
     "tags": []
    },
@@ -483,7 +483,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "22c8c0fe",
+   "id": "6920e975",
    "metadata": {
     "tags": [
      "solution"
@@ -503,15 +503,17 @@
     "    x, target=y, baselines=blurred_baselines\n",
     ")\n",
     "\n",
+    "attributions_blurred = attributions_blurred.cpu().numpy()\n",
+    "\n",
     "# Plotting\n",
-    "for attr, im, lbl in zip(attributions, x.cpu().numpy(), y.cpu().numpy()):\n",
+    "for attr, im, lbl in zip(attributions_blurred, x.cpu().numpy(), y.cpu().numpy()):\n",
     "    print(f\"Class {lbl}\")\n",
     "    visualize_color_attribution(attr, im)"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "220cc456",
+   "id": "de810272",
    "metadata": {
     "tags": []
    },
@@ -527,7 +529,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "83bbf9b4",
+   "id": "7a6f89dc",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-info\"><h2>BONUS Task: Using different attributions.</h2>\n",
@@ -541,11 +543,11 @@
   },
   {
    "cell_type": "markdown",
-   "id": "69d66239",
+   "id": "f3e293c5",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-success\"><h2>Checkpoint 2</h2>\n",
-    "Let us know on the exercise chat when you've reached this point!\n",
+    "Put up your green sticky note when you've reached this point!\n",
     "\n",
     "At this point we have:\n",
     "\n",
@@ -561,7 +563,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0d20d9cf",
+   "id": "fb3048e1",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -574,9 +576,9 @@
     "**What is a counterfactual?**\n",
     "\n",
     "You've learned about adversarial examples in the lecture on failure modes. These are the imperceptible or noisy changes to an image that drastically changes a classifier's opinion.\n",
-    "Counterfactual explanations are the useful cousins of adversarial examples. They are *perceptible* and *informative* changes to an image that changes a classifier's opinion.\n",
+    "Counterfactual explanations are the useful cousins of adversarial examples. They are *perceptible* and *informative* changes to an image that change a classifier's opinion.\n",
     "\n",
-    "In the image below you can see the difference between the two. In the first column are MNIST images along with their classifictaions, and in the second column are counterfactual explanations to *change* that class. You can see that in both cases a human being would (hopefully) agree with the new classification. By comparing the two columns, we can therefore begin to define what makes each digit special.\n",
+    "In the image below you can see the difference between the two. In the first column are (non-color) MNIST images along with their classifications, and in the second column are counterfactual explanations to *change* that class. You can see that in both cases a human being would (hopefully) agree with the new classification. By comparing the two columns, we can therefore begin to define what makes each digit special.\n",
     "\n",
     "In contrast, the third and fourth columns show an MNIST image and a corresponding adversarial example. Here the network returns a prediction that most human beings (who aren't being facetious) would strongly disagree with.\n",
     "\n",
@@ -589,7 +591,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "18cda8e3",
+   "id": "88f77592",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -612,7 +614,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a4ea7f03",
+   "id": "6742eca7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -644,7 +646,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "53e15fe2",
+   "id": "3ed703e3",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -659,7 +661,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "33e1c7fa",
+   "id": "981023f6",
    "metadata": {
     "tags": [
      "solution"
@@ -676,14 +678,14 @@
   },
   {
    "cell_type": "markdown",
-   "id": "49a9a89d",
+   "id": "fdf97172",
    "metadata": {
     "tags": []
    },
    "source": [
     "<div class=\"alert alert-block alert-warning\"><h3>Hyper-parameter choices</h3>\n",
     "<ul>\n",
-    "<li>Are any of the hyperparameters you choose above constrained in some way?</li>\n",
+    "<li>Are any of the hyperparameters above constrained in some way?</li>\n",
     "<li>What would happen if you chose a depth of 10 for the UNet?</li>\n",
     "<li>Is there a minimum size for the style space? Why or why not?</li>\n",
     "</ul>"
@@ -691,7 +693,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "fb9b4a59",
+   "id": "bcb08a78",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -708,7 +710,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "326f1c50",
+   "id": "fa38c00f",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": [
@@ -722,7 +724,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ef8dd1e5",
+   "id": "fd541400",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -733,7 +735,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "069f2804",
+   "id": "1e91e416",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -743,7 +745,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "380256c9",
+   "id": "b0377825",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -761,7 +763,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "576200cc",
+   "id": "f4939a8d",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -773,7 +775,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2885a3a0",
+   "id": "9a6ca642",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -792,7 +794,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9cc5eb23",
+   "id": "382fa57e",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -801,7 +803,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5eb04071",
+   "id": "ccf3b275",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -817,7 +819,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "805dfbf2",
+   "id": "0a424199",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -826,7 +828,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "8fab1a3e",
+   "id": "fe335f99",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -838,7 +840,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "dc899114",
+   "id": "977ce4c8",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -851,7 +853,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7e13436e",
+   "id": "7bf8cb24",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -865,7 +867,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1cb60fb1",
+   "id": "0d5ee1e1",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -877,7 +879,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e8441eda",
+   "id": "43c93243",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -897,7 +899,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6bb72378",
+   "id": "9a789ba4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -921,7 +923,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "af5d018a",
+   "id": "91b15866",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -931,7 +933,26 @@
   },
   {
    "cell_type": "markdown",
-   "id": "97aaeef6",
+   "id": "87d5a8a1",
+   "metadata": {},
+   "source": [
+    "<div class=\"alert alert-block alert-success\"><h2>Checkpoint 3</h2>\n",
+    "Put up your green sticky note when you've reached this point!\n",
+    "\n",
+    "At this point we have:\n",
+    "\n",
+    "- Loaded a classifier that classifies MNIST-like images by color, but we don't know how!\n",
+    "- Tried applying Integrated Gradients to find out what the classifier is looking at - with little success.\n",
+    "- Discovered the effect of changing the baseline on the output of integrated gradients.\n",
+    "- Defined the hyperparameters for a StarGAN to create counterfactual images.\n",
+    "\n",
+    "Next up, we will define the training loop for the StarGAN.\n",
+    "</div>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3a1e7dd9",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -953,7 +974,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "74570a78",
+   "id": "c6348be7",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -965,7 +986,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e53b9f64",
+   "id": "e67449b2",
    "metadata": {
     "lines_to_next_cell": 2,
     "tags": [
@@ -1035,7 +1056,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "41e0e738",
+   "id": "f6cd135b",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1047,7 +1068,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b49fd330",
+   "id": "cdecbbbe",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1063,7 +1084,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "77d4f778",
+   "id": "2c7d13e4",
    "metadata": {
     "tags": []
    },
@@ -1078,7 +1099,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2258df98",
+   "id": "689db81d",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1090,7 +1111,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "fa248f8b",
+   "id": "c3f514fa",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1112,7 +1133,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c2db79f3",
+   "id": "27ee28e5",
    "metadata": {
     "tags": []
    },
@@ -1122,13 +1143,13 @@
     "The same method can be used to create a StarGAN with different basic elements.\n",
     "For example, you can change the archictecture of the generators, or of the discriminator to better fit your data in the future.\n",
     "\n",
-    "You know the drill... let us know on the exercise chat when you have arrived here!\n",
+    "You know the drill... put up your green sticky note when you have arrived here!\n",
     "</div>"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "de5cb982",
+   "id": "e1c592af",
    "metadata": {
     "tags": []
    },
@@ -1138,7 +1159,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ac566b2e",
+   "id": "d2374be1",
    "metadata": {
     "tags": []
    },
@@ -1155,7 +1176,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "af5bbb66",
+   "id": "c038ee90",
    "metadata": {
     "title": "Loading the test dataset"
    },
@@ -1175,7 +1196,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e5392251",
+   "id": "2dd7e283",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1187,7 +1208,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "02808bdf",
+   "id": "1910ed0f",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1200,7 +1221,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e2e7f289",
+   "id": "99188e3e",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1210,7 +1231,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "9b936567",
+   "id": "fe2b4929",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1228,7 +1249,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "caa3e4ec",
+   "id": "99768fc9",
    "metadata": {
     "tags": [
      "solution"
@@ -1265,7 +1286,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "35dd9913",
+   "id": "ef56e941",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1277,7 +1298,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "13fa9bfa",
+   "id": "1612f8ce",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1290,7 +1311,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "449428b9",
+   "id": "2239beb9",
    "metadata": {
     "tags": []
    },
@@ -1305,7 +1326,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "cb3d6985",
+   "id": "168352ff",
    "metadata": {
     "tags": []
    },
@@ -1316,7 +1337,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3096db55",
+   "id": "f24cdc48",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1330,7 +1351,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "19fd90a6",
+   "id": "b093ffc8",
    "metadata": {
     "tags": []
    },
@@ -1345,7 +1366,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e0967ec3",
+   "id": "be5c5a98",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1360,7 +1381,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b3312425",
+   "id": "0891fcf7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1381,7 +1402,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "42c617e2",
+   "id": "b080dfe1",
    "metadata": {
     "title": "Another visualization function"
    },
@@ -1410,7 +1431,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3075be99",
+   "id": "e41e00c9",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1426,7 +1447,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "d70f3ff4",
+   "id": "ac7678a7",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1442,7 +1463,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "593ba2db",
+   "id": "3e06986c",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1457,7 +1478,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "863c8dca",
+   "id": "7c9a99c4",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1471,7 +1492,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "602db341",
+   "id": "b1affa61",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1491,7 +1512,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "cf1e78db",
+   "id": "4bf0e982",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1527,7 +1548,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6e37fd94",
+   "id": "7f8f6ca5",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1541,7 +1562,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5b55e69e",
+   "id": "537edaa0",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1557,7 +1578,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3862ab46",
+   "id": "c4dd727d",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1573,7 +1594,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "bedeaf5c",
+   "id": "c5a5e2c0",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1596,7 +1617,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "75e28b2c",
+   "id": "a9ede52b",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-info\"><h3>Task 5.1: Explore the style space</h3>\n",
@@ -1608,7 +1629,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d2ffe8e5",
+   "id": "af4a9d95",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1643,7 +1664,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "559a66a2",
+   "id": "be742cb1",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1659,7 +1680,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4c268ad8",
+   "id": "09032a46",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1686,7 +1707,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7e9d9184",
+   "id": "4d9a7277",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1700,7 +1721,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "45005d21",
+   "id": "3016bf5c",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1717,7 +1738,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b6e19ec8",
+   "id": "b47bc9a2",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1740,7 +1761,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7aba6e43",
+   "id": "5434a70a",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1749,7 +1770,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c90115e2",
+   "id": "9d066df4",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-warning\"><h3>Questions</h3>\n",
@@ -1761,7 +1782,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6dd82862",
+   "id": "e4295029",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-success\"><h2>Checkpoint 5</h2>\n",
@@ -1779,7 +1800,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b6389689",
+   "id": "3cd7976f",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1801,7 +1822,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "764a648d",
+   "id": "ffeee5b2",
    "metadata": {
     "tags": [
      "solution"