diff --git a/code/testing_code.ipynb b/code/testing_code.ipynb
index db3fff11..efa99bd6 100644
--- a/code/testing_code.ipynb
+++ b/code/testing_code.ipynb
@@ -23,6 +23,19 @@
         "Testing functions and classes"
       ]
     },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "biaboE9QGiiI",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "from torchvision import datasets"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
     {
       "cell_type": "code",
       "metadata": {
@@ -32,7 +45,8 @@
       },
       "source": [
         "# Serve per il mapping fra label sul file con le labels per canzone di test e l'indice della label che la rete considera.\n",
-        "LABELS_FROM_FILE = {'angry' : 0, 'calming' : 1, 'happy' : 2, 'normal' : 3, 'sad' : 4}"
+        "LABELS_FROM_FILE = {'angry' : 0, 'calming' : 1, 'happy' : 2, 'normal' : 3, 'sad' : 4}\n",
+        "NUM_CLASSES = 5"
       ],
       "execution_count": 0,
       "outputs": []