diff --git a/exercise.ipynb b/exercise.ipynb
index 50d5adb..4d5da75 100644
--- a/exercise.ipynb
+++ b/exercise.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "be4e7a97",
+   "id": "5836c5fd",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -29,7 +29,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a1e6c3cd",
+   "id": "c7ee63fa",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -41,7 +41,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "bc9d8f39",
+   "id": "25948d26",
    "metadata": {},
    "source": [
     "\n",
@@ -54,7 +54,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5a28affe",
+   "id": "8b8c18b6",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -68,7 +68,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3fb92e01",
+   "id": "2d9a1fa2",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -84,7 +84,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "225052d7",
+   "id": "33bedaea",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -102,7 +102,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a878ce13",
+   "id": "d8c857dc",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -113,7 +113,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "1f8d5401",
+   "id": "04aa9880",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -130,7 +130,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "19f50056",
+   "id": "8d3cf42c",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": [
@@ -155,7 +155,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6dd3913c",
+   "id": "252e596c",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -166,7 +166,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "060a05f5",
+   "id": "2bb3f66f",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -195,7 +195,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2616d093",
+   "id": "ad6e8d02",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -212,7 +212,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "76fae027",
+   "id": "8ff2253c",
    "metadata": {},
    "source": [
     "# Part 2: Using Integrated Gradients to find what the classifier knows\n",
@@ -222,7 +222,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5284eaf7",
+   "id": "137f5e29",
    "metadata": {},
    "source": [
     "## Attributions through integrated gradients\n",
@@ -235,7 +235,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "57eb7e30",
+   "id": "877b0dc5",
    "metadata": {
     "tags": []
    },
@@ -253,7 +253,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5dd749a7",
+   "id": "346433e9",
    "metadata": {
     "tags": []
    },
@@ -269,7 +269,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7d56e7b9",
+   "id": "2ed39b81",
    "metadata": {
     "tags": [
      "task"
@@ -290,7 +290,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f1942f36",
+   "id": "d182aa7a",
    "metadata": {
     "tags": []
    },
@@ -303,7 +303,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7e5d1815",
+   "id": "ef1e380d",
    "metadata": {
     "lines_to_next_cell": 2,
     "tags": []
@@ -315,7 +315,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d5103ffb",
+   "id": "3fce7d60",
    "metadata": {
     "tags": []
    },
@@ -343,7 +343,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "858a7c14",
+   "id": "d38da2c2",
    "metadata": {
     "tags": []
    },
@@ -356,7 +356,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "78c186ef",
+   "id": "aff26564",
    "metadata": {
     "lines_to_next_cell": 2
    },
@@ -370,7 +370,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5e3aa105",
+   "id": "53ffc390",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -383,7 +383,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8171ef5a",
+   "id": "160b27d5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -408,7 +408,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0b65a750",
+   "id": "42462cab",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -422,7 +422,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "06d9fccd",
+   "id": "b4cd9eac",
    "metadata": {},
    "source": [
     "\n",
@@ -448,7 +448,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "1cf3195f",
+   "id": "dc603cd7",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-info\"><h4>Task 2.3: Use random noise as a baseline</h4>\n",
@@ -460,7 +460,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "fbefdb10",
+   "id": "9be598b0",
    "metadata": {
     "tags": [
      "task"
@@ -482,7 +482,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e31bd3cf",
+   "id": "52c0979f",
    "metadata": {
     "tags": []
    },
@@ -496,7 +496,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "98de5d82",
+   "id": "d392024c",
    "metadata": {
     "tags": [
      "task"
@@ -520,7 +520,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "290bdbf3",
+   "id": "30b56566",
    "metadata": {
     "tags": []
    },
@@ -536,7 +536,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "659c5758",
+   "id": "ee7c164c",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-info\"><h2>BONUS Task: Using different attributions.</h2>\n",
@@ -550,7 +550,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b93071a2",
+   "id": "71b2e6fa",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-success\"><h2>Checkpoint 2</h2>\n",
@@ -570,7 +570,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e9f2b7ae",
+   "id": "e0966908",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -598,7 +598,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "92223fdc",
+   "id": "50bfe53d",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -621,7 +621,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5a72db0b",
+   "id": "b1ab5f6d",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -653,7 +653,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "01335c40",
+   "id": "7029ce21",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -668,7 +668,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8d0f0203",
+   "id": "48bc8945",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": [
@@ -689,7 +689,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "186f40a5",
+   "id": "156d8774",
    "metadata": {
     "tags": []
    },
@@ -704,7 +704,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "af4cf127",
+   "id": "48a541a1",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -721,7 +721,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "94f07b70",
+   "id": "433257a6",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": [
@@ -735,7 +735,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c2a072a4",
+   "id": "b16b319e",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -746,7 +746,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f7bb217b",
+   "id": "1e85b7b2",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -756,7 +756,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "dfdd1272",
+   "id": "dc1c55cf",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -774,7 +774,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7ddd581a",
+   "id": "33e0ea83",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -786,7 +786,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2f7fe43c",
+   "id": "6bbe594d",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -805,7 +805,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f01095a7",
+   "id": "7b5f7065",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -814,7 +814,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "bad091aa",
+   "id": "c4886cf1",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -830,7 +830,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e49c55e0",
+   "id": "1f2d9558",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -839,7 +839,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "af990b9a",
+   "id": "6632b7c6",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -851,7 +851,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "291f2062",
+   "id": "2b336a53",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -864,7 +864,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "919aea73",
+   "id": "b7c49115",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -878,7 +878,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "26f2a9c1",
+   "id": "e076c809",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -890,7 +890,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "54780ac3",
+   "id": "9d7003f8",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -910,7 +910,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e7187b29",
+   "id": "f3b60479",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -934,7 +934,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "78521f46",
+   "id": "103b94b0",
    "metadata": {
     "lines_to_next_cell": 2
    },
@@ -946,7 +946,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "036e6086",
+   "id": "d146e4f8",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -968,7 +968,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3b756b7a",
+   "id": "11e0f5dd",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -980,7 +980,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "20847a1b",
+   "id": "79722273",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": [
@@ -1091,7 +1091,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "78403aaa",
+   "id": "5a0870e3",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1103,7 +1103,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9bb6c797",
+   "id": "9c9a471f",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1119,7 +1119,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c32cbd8b",
+   "id": "f01e0c51",
    "metadata": {
     "tags": []
    },
@@ -1134,7 +1134,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "4c86dd42",
+   "id": "3c545933",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1146,7 +1146,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8f728ded",
+   "id": "74225676",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1168,7 +1168,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "028d3bbc",
+   "id": "b83dfcd2",
    "metadata": {
     "tags": []
    },
@@ -1184,7 +1184,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b63aac86",
+   "id": "740be1a4",
    "metadata": {
     "tags": []
    },
@@ -1194,7 +1194,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b0ad5935",
+   "id": "2a44a05e",
    "metadata": {
     "tags": []
    },
@@ -1211,7 +1211,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0cf2b9a5",
+   "id": "8d711fbe",
    "metadata": {
     "title": "Loading the test dataset"
    },
@@ -1231,7 +1231,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "35cc9b35",
+   "id": "7c1d56b5",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1243,7 +1243,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4146969c",
+   "id": "f649373f",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1256,7 +1256,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f99f676c",
+   "id": "4a3c2e42",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1266,7 +1266,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a7450339",
+   "id": "746827ec",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1284,7 +1284,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "58973345",
+   "id": "855b63d4",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": [
@@ -1320,7 +1320,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "20af6915",
+   "id": "9e21d254",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1332,7 +1332,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4543762f",
+   "id": "143cb6c6",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1345,7 +1345,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "d0b8a2ec",
+   "id": "a15cdb80",
    "metadata": {
     "tags": []
    },
@@ -1360,7 +1360,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "61e58d7f",
+   "id": "2d2ced5b",
    "metadata": {
     "tags": []
    },
@@ -1371,7 +1371,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b8dc4640",
+   "id": "24639718",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1385,7 +1385,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ed9b7104",
+   "id": "8afab4b7",
    "metadata": {
     "tags": []
    },
@@ -1400,7 +1400,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0add70ab",
+   "id": "927a2361",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1415,7 +1415,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "274d8226",
+   "id": "87aaa903",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1436,7 +1436,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b6b54bbe",
+   "id": "3633b841",
    "metadata": {
     "title": "Another visualization function"
    },
@@ -1465,7 +1465,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "979ca20f",
+   "id": "32f5d3ba",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1481,7 +1481,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b779b535",
+   "id": "89ecd3a9",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1497,7 +1497,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c44f0cf1",
+   "id": "19612f27",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1512,7 +1512,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f402bb74",
+   "id": "f63764f2",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1526,7 +1526,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "12542a6d",
+   "id": "1bcaa03b",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1546,7 +1546,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8d3e6d86",
+   "id": "d3b3324a",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1582,7 +1582,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ca1d1242",
+   "id": "81ea3fdd",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1596,7 +1596,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6f530e82",
+   "id": "92c0dba0",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1612,7 +1612,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5904b20b",
+   "id": "93ad40a6",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1628,7 +1628,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "72ca6d87",
+   "id": "884f4f82",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1651,7 +1651,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a0baa037",
+   "id": "344c70a0",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-info\"><h3>Task 5.1: Explore the style space</h3>\n",
@@ -1663,7 +1663,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "bc0062db",
+   "id": "bc06e146",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1698,7 +1698,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ba428131",
+   "id": "518c6ddb",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1714,7 +1714,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2ee2e061",
+   "id": "40fa1f61",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1741,7 +1741,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a26903b8",
+   "id": "a204c00f",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1755,7 +1755,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7f72dfb5",
+   "id": "16971e9c",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1772,7 +1772,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "be62a09b",
+   "id": "5d04746b",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1794,7 +1794,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f16385e4",
+   "id": "7dd5245c",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-warning\"><h3>Questions</h3>\n",
@@ -1806,7 +1806,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e0f59eee",
+   "id": "4ab3793f",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-success\"><h2>Checkpoint 5</h2>\n",
@@ -1824,7 +1824,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0bbc6cd2",
+   "id": "2d4e3c2a",
    "metadata": {},
    "source": [
     "# Bonus!\n",
@@ -1834,46 +1834,7 @@
     "<li> What happens if you don't use the EMA model? </li>\n",
     "<li> What happens if you change the learning rates? </li>\n",
     "<li> What happens if you add a Sigmoid activation to the output of the style encoder? </li>\n",
-    "See what else you can think of, and see how finnicky training a GAN can be!\n",
-    "\n",
-    "# %% [markdown] tags=[\"solution\"]\n",
-    "The colors for the classes are sampled from matplotlib colormaps! They are the four seasons: spring, summer, autumn, and winter.\n",
-    "Check your style space again to see if you can see the patterns now!\n",
-    "\n",
-    "# %% tags=[\"solution\"]\n",
-    "Let's plot the colormaps\n",
-    "import matplotlib as mpl\n",
-    "import numpy as np\n",
-    "\n",
-    "\n",
-    "def plot_color_gradients(cmap_list):\n",
-    "    gradient = np.linspace(0, 1, 256)\n",
-    "    gradient = np.vstack((gradient, gradient))\n",
-    "\n",
-    "    # Create figure and adjust figure height to number of colormaps\n",
-    "    nrows = len(cmap_list)\n",
-    "    figh = 0.35 + 0.15 + (nrows + (nrows - 1) * 0.1) * 0.22\n",
-    "    fig, axs = plt.subplots(nrows=nrows + 1, figsize=(6.4, figh))\n",
-    "    fig.subplots_adjust(top=1 - 0.35 / figh, bottom=0.15 / figh, left=0.2, right=0.99)\n",
-    "\n",
-    "    for ax, name in zip(axs, cmap_list):\n",
-    "        ax.imshow(gradient, aspect=\"auto\", cmap=mpl.colormaps[name])\n",
-    "        ax.text(\n",
-    "            -0.01,\n",
-    "            0.5,\n",
-    "            name,\n",
-    "            va=\"center\",\n",
-    "            ha=\"right\",\n",
-    "            fontsize=10,\n",
-    "            transform=ax.transAxes,\n",
-    "        )\n",
-    "\n",
-    "    # Turn off *all* ticks & spines, not just the ones with colormaps.\n",
-    "    for ax in axs:\n",
-    "        ax.set_axis_off()\n",
-    "\n",
-    "\n",
-    "plot_color_gradients([\"spring\", \"summer\", \"autumn\", \"winter\"])"
+    "See what else you can think of, and see how finnicky training a GAN can be!"
    ]
   }
  ],
diff --git a/solution.ipynb b/solution.ipynb
index d4808b6..2fbb96f 100644
--- a/solution.ipynb
+++ b/solution.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "be4e7a97",
+   "id": "5836c5fd",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -29,7 +29,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a1e6c3cd",
+   "id": "c7ee63fa",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -41,7 +41,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "bc9d8f39",
+   "id": "25948d26",
    "metadata": {},
    "source": [
     "\n",
@@ -54,7 +54,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5a28affe",
+   "id": "8b8c18b6",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -68,7 +68,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3fb92e01",
+   "id": "2d9a1fa2",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -84,7 +84,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "225052d7",
+   "id": "33bedaea",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -102,7 +102,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a878ce13",
+   "id": "d8c857dc",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -113,7 +113,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "1f8d5401",
+   "id": "04aa9880",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -130,7 +130,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "dc1f9832",
+   "id": "c15b2288",
    "metadata": {
     "tags": [
      "solution"
@@ -154,7 +154,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6dd3913c",
+   "id": "252e596c",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -165,7 +165,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "060a05f5",
+   "id": "2bb3f66f",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -194,7 +194,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2616d093",
+   "id": "ad6e8d02",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -211,7 +211,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "76fae027",
+   "id": "8ff2253c",
    "metadata": {},
    "source": [
     "# Part 2: Using Integrated Gradients to find what the classifier knows\n",
@@ -221,7 +221,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5284eaf7",
+   "id": "137f5e29",
    "metadata": {},
    "source": [
     "## Attributions through integrated gradients\n",
@@ -234,7 +234,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "57eb7e30",
+   "id": "877b0dc5",
    "metadata": {
     "tags": []
    },
@@ -252,7 +252,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5dd749a7",
+   "id": "346433e9",
    "metadata": {
     "tags": []
    },
@@ -268,7 +268,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9e77b469",
+   "id": "e5cd638a",
    "metadata": {
     "tags": [
      "solution"
@@ -292,7 +292,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f1942f36",
+   "id": "d182aa7a",
    "metadata": {
     "tags": []
    },
@@ -305,7 +305,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7e5d1815",
+   "id": "ef1e380d",
    "metadata": {
     "lines_to_next_cell": 2,
     "tags": []
@@ -317,7 +317,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d5103ffb",
+   "id": "3fce7d60",
    "metadata": {
     "tags": []
    },
@@ -345,7 +345,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "858a7c14",
+   "id": "d38da2c2",
    "metadata": {
     "tags": []
    },
@@ -358,7 +358,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "78c186ef",
+   "id": "aff26564",
    "metadata": {
     "lines_to_next_cell": 2
    },
@@ -372,7 +372,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5e3aa105",
+   "id": "53ffc390",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -385,7 +385,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8171ef5a",
+   "id": "160b27d5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -410,7 +410,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0b65a750",
+   "id": "42462cab",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -424,7 +424,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "06d9fccd",
+   "id": "b4cd9eac",
    "metadata": {},
    "source": [
     "\n",
@@ -450,7 +450,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "1cf3195f",
+   "id": "dc603cd7",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-info\"><h4>Task 2.3: Use random noise as a baseline</h4>\n",
@@ -462,7 +462,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b09d2144",
+   "id": "768a9e3b",
    "metadata": {
     "tags": [
      "solution"
@@ -488,7 +488,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e31bd3cf",
+   "id": "52c0979f",
    "metadata": {
     "tags": []
    },
@@ -502,7 +502,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "af1e6318",
+   "id": "38a82581",
    "metadata": {
     "tags": [
      "solution"
@@ -532,7 +532,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "290bdbf3",
+   "id": "30b56566",
    "metadata": {
     "tags": []
    },
@@ -548,7 +548,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "659c5758",
+   "id": "ee7c164c",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-info\"><h2>BONUS Task: Using different attributions.</h2>\n",
@@ -562,7 +562,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b93071a2",
+   "id": "71b2e6fa",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-success\"><h2>Checkpoint 2</h2>\n",
@@ -582,7 +582,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e9f2b7ae",
+   "id": "e0966908",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -610,7 +610,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "92223fdc",
+   "id": "50bfe53d",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -633,7 +633,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5a72db0b",
+   "id": "b1ab5f6d",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -665,7 +665,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "01335c40",
+   "id": "7029ce21",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -680,7 +680,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d86d3dcc",
+   "id": "8e83c444",
    "metadata": {
     "tags": [
      "solution"
@@ -697,7 +697,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "186f40a5",
+   "id": "156d8774",
    "metadata": {
     "tags": []
    },
@@ -712,7 +712,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "af4cf127",
+   "id": "48a541a1",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -729,7 +729,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1f145891",
+   "id": "1608ecf2",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": [
@@ -743,7 +743,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c2a072a4",
+   "id": "b16b319e",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -754,7 +754,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f7bb217b",
+   "id": "1e85b7b2",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -764,7 +764,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "dfdd1272",
+   "id": "dc1c55cf",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -782,7 +782,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7ddd581a",
+   "id": "33e0ea83",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -794,7 +794,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2f7fe43c",
+   "id": "6bbe594d",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -813,7 +813,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f01095a7",
+   "id": "7b5f7065",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -822,7 +822,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "bad091aa",
+   "id": "c4886cf1",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -838,7 +838,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e49c55e0",
+   "id": "1f2d9558",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -847,7 +847,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "af990b9a",
+   "id": "6632b7c6",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -859,7 +859,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "291f2062",
+   "id": "2b336a53",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -872,7 +872,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "919aea73",
+   "id": "b7c49115",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -886,7 +886,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "26f2a9c1",
+   "id": "e076c809",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -898,7 +898,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "54780ac3",
+   "id": "9d7003f8",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -918,7 +918,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e7187b29",
+   "id": "f3b60479",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -942,7 +942,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "78521f46",
+   "id": "103b94b0",
    "metadata": {
     "lines_to_next_cell": 2
    },
@@ -954,7 +954,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "036e6086",
+   "id": "d146e4f8",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -976,7 +976,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3b756b7a",
+   "id": "11e0f5dd",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -988,7 +988,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8256f4b5",
+   "id": "6fefc28d",
    "metadata": {
     "lines_to_next_cell": 2,
     "tags": [
@@ -1058,7 +1058,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "78403aaa",
+   "id": "5a0870e3",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1070,7 +1070,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9bb6c797",
+   "id": "9c9a471f",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1086,7 +1086,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c32cbd8b",
+   "id": "f01e0c51",
    "metadata": {
     "tags": []
    },
@@ -1101,7 +1101,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "4c86dd42",
+   "id": "3c545933",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1113,7 +1113,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8f728ded",
+   "id": "74225676",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1135,7 +1135,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "028d3bbc",
+   "id": "b83dfcd2",
    "metadata": {
     "tags": []
    },
@@ -1151,7 +1151,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b63aac86",
+   "id": "740be1a4",
    "metadata": {
     "tags": []
    },
@@ -1161,7 +1161,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b0ad5935",
+   "id": "2a44a05e",
    "metadata": {
     "tags": []
    },
@@ -1178,7 +1178,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0cf2b9a5",
+   "id": "8d711fbe",
    "metadata": {
     "title": "Loading the test dataset"
    },
@@ -1198,7 +1198,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "35cc9b35",
+   "id": "7c1d56b5",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1210,7 +1210,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4146969c",
+   "id": "f649373f",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1223,7 +1223,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f99f676c",
+   "id": "4a3c2e42",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1233,7 +1233,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a7450339",
+   "id": "746827ec",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1251,7 +1251,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f2bfd025",
+   "id": "870dd94a",
    "metadata": {
     "tags": [
      "solution"
@@ -1288,7 +1288,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "20af6915",
+   "id": "9e21d254",
    "metadata": {
     "lines_to_next_cell": 0,
     "tags": []
@@ -1300,7 +1300,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4543762f",
+   "id": "143cb6c6",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1313,7 +1313,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "d0b8a2ec",
+   "id": "a15cdb80",
    "metadata": {
     "tags": []
    },
@@ -1328,7 +1328,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "61e58d7f",
+   "id": "2d2ced5b",
    "metadata": {
     "tags": []
    },
@@ -1339,7 +1339,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b8dc4640",
+   "id": "24639718",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1353,7 +1353,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ed9b7104",
+   "id": "8afab4b7",
    "metadata": {
     "tags": []
    },
@@ -1368,7 +1368,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0add70ab",
+   "id": "927a2361",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1383,7 +1383,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "274d8226",
+   "id": "87aaa903",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1404,7 +1404,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b6b54bbe",
+   "id": "3633b841",
    "metadata": {
     "title": "Another visualization function"
    },
@@ -1433,7 +1433,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "979ca20f",
+   "id": "32f5d3ba",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1449,7 +1449,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b779b535",
+   "id": "89ecd3a9",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1465,7 +1465,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c44f0cf1",
+   "id": "19612f27",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1480,7 +1480,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f402bb74",
+   "id": "f63764f2",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1494,7 +1494,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "12542a6d",
+   "id": "1bcaa03b",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1514,7 +1514,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8d3e6d86",
+   "id": "d3b3324a",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1550,7 +1550,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ca1d1242",
+   "id": "81ea3fdd",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1564,7 +1564,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6f530e82",
+   "id": "92c0dba0",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1580,7 +1580,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5904b20b",
+   "id": "93ad40a6",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1596,7 +1596,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "72ca6d87",
+   "id": "884f4f82",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1619,7 +1619,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a0baa037",
+   "id": "344c70a0",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-info\"><h3>Task 5.1: Explore the style space</h3>\n",
@@ -1631,7 +1631,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "bc0062db",
+   "id": "bc06e146",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1666,7 +1666,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ba428131",
+   "id": "518c6ddb",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1682,7 +1682,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2ee2e061",
+   "id": "40fa1f61",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1709,7 +1709,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a26903b8",
+   "id": "a204c00f",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1723,7 +1723,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7f72dfb5",
+   "id": "16971e9c",
    "metadata": {
     "lines_to_next_cell": 0
    },
@@ -1740,7 +1740,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "be62a09b",
+   "id": "5d04746b",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1762,7 +1762,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f16385e4",
+   "id": "7dd5245c",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-warning\"><h3>Questions</h3>\n",
@@ -1774,7 +1774,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e0f59eee",
+   "id": "4ab3793f",
    "metadata": {},
    "source": [
     "<div class=\"alert alert-block alert-success\"><h2>Checkpoint 5</h2>\n",
@@ -1792,7 +1792,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0bbc6cd2",
+   "id": "2d4e3c2a",
    "metadata": {},
    "source": [
     "# Bonus!\n",
@@ -1802,14 +1802,34 @@
     "<li> What happens if you don't use the EMA model? </li>\n",
     "<li> What happens if you change the learning rates? </li>\n",
     "<li> What happens if you add a Sigmoid activation to the output of the style encoder? </li>\n",
-    "See what else you can think of, and see how finnicky training a GAN can be!\n",
-    "\n",
-    "# %% [markdown] tags=[\"solution\"]\n",
+    "See what else you can think of, and see how finnicky training a GAN can be!"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1cd1cd8b",
+   "metadata": {
+    "tags": [
+     "solution"
+    ]
+   },
+   "source": [
     "The colors for the classes are sampled from matplotlib colormaps! They are the four seasons: spring, summer, autumn, and winter.\n",
-    "Check your style space again to see if you can see the patterns now!\n",
-    "\n",
-    "# %% tags=[\"solution\"]\n",
-    "Let's plot the colormaps\n",
+    "Check your style space again to see if you can see the patterns now!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9ef23fc4",
+   "metadata": {
+    "tags": [
+     "solution"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "# Let's plot the colormaps\n",
     "import matplotlib as mpl\n",
     "import numpy as np\n",
     "\n",