diff --git a/MC/MC Prediction Solution.ipynb b/MC/MC Prediction Solution.ipynb
index 25da5f3ca..d72467094 100644
--- a/MC/MC Prediction Solution.ipynb	
+++ b/MC/MC Prediction Solution.ipynb	
@@ -92,9 +92,9 @@
     "        states_in_episode = set([tuple(x[0]) for x in episode])\n",
     "        for state in states_in_episode:\n",
     "            # Find the first occurance of the state in the episode\n",
-    "            first_occurence_idx = next(i for i,x in enumerate(episode) if x[0] == state)\n",
+    "            first_occurence_idx = [i for i,x in enumerate(episode) if x[0] == state]\n",
     "            # Sum up all rewards since the first occurance\n",
-    "            G = sum([x[2]*(discount_factor**i) for i,x in enumerate(episode[first_occurence_idx:])])\n",
+    "            G = sum([x[2]*(discount_factor**i) for i,x in enumerate(episode[first_occurence_idx[0]:])])\n",
     "            # Calculate average return for this state over all sampled episodes\n",
     "            returns_sum[state] += G\n",
     "            returns_count[state] += 1.0\n",