diff --git a/forrest-test.ipynb b/forrest-test.ipynb
index 7487af5037c3abbfdcd9f113803a6d1aba414b0c..5d52705e4da89caf0fb4f38c7ad9fb2fe8184adb 100644
--- a/forrest-test.ipynb
+++ b/forrest-test.ipynb
@@ -24,7 +24,7 @@
     "from machineLearning.rf import (\n",
     "    RandomForest, DecisionTree,\n",
     "    Gini, Entropy, MAE, MSE,\n",
-    "    Mode, Mean, Confidence,\n",
+    "    Mode, Mean, Confidence, Probabilities,\n",
     "    CART, ID3, C45,\n",
     "    AdaBoosting, GradientBoosting,\n",
     "    Majority, Confidence, Average, Median\n",
@@ -85,7 +85,8 @@
     "# Shuffle the training data\n",
     "trainIndex = np.random.permutation(len(trainData))\n",
     "trainData = trainData[trainIndex]\n",
-    "trainLabels = trainLabels[trainIndex]"
+    "trainLabels = trainLabels[trainIndex]\n",
+    "trainTargets = trainTargets[trainIndex]"
    ]
   },
   {