diff --git a/decision_tree.ipynb b/decision_tree.ipynb
index 07da773..ed38446 100644
--- a/decision_tree.ipynb
+++ b/decision_tree.ipynb
@@ -54,7 +54,9 @@
    "source": [
     "# Column labels.\n",
     "# These are used only to print the tree.\n",
-    "header = [\"color\", \"diameter\", \"label\"]"
+    "header = [\"color\", \"diameter\", \"label\"]\n",
+    "label = [row[-1] for row in training_data]"
+    
    ]
   },
   {