diff --git a/tv-script-generation/dlnd_tv_script_generation.ipynb b/tv-script-generation/dlnd_tv_script_generation.ipynb
index de7b75f..ac665fa 100644
--- a/tv-script-generation/dlnd_tv_script_generation.ipynb
+++ b/tv-script-generation/dlnd_tv_script_generation.ipynb
@@ -1992,6 +1992,58 @@
     "tests.test_get_tensors(get_tensors)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 511,
+   "metadata": {
+    "collapsed": false,
+    "scrolled": true
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(3,)\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "3"
+      ]
+     },
+     "execution_count": 511,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "t = np.array([1,2,3])\n",
+    "print(t.shape)\n",
+    "np.extract(t==3, t)[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 458,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "0"
+      ]
+     },
+     "execution_count": 458,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": []
+  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -2005,7 +2057,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 274,
+   "execution_count": 472,
    "metadata": {
     "collapsed": false,
     "deletable": true,
@@ -2016,7 +2068,20 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Tests Passed\n"
+      "0.8\n"
+     ]
+    },
+    {
+     "ename": "IndexError",
+     "evalue": "only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
+      "\u001b[0;32m<ipython-input-472-ec56654dfaa2>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[0mDON\u001b[0m\u001b[0;31m'\u001b[0m\u001b[0mT\u001b[0m \u001b[0mMODIFY\u001b[0m \u001b[0mANYTHING\u001b[0m \u001b[0mIN\u001b[0m \u001b[0mTHIS\u001b[0m \u001b[0mCELL\u001b[0m \u001b[0mTHAT\u001b[0m \u001b[0mIS\u001b[0m \u001b[0mBELOW\u001b[0m \u001b[0mTHIS\u001b[0m \u001b[0mLINE\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m \"\"\"\n\u001b[0;32m---> 18\u001b[0;31m \u001b[0mtests\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtest_pick_word\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpick_word\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
+      "\u001b[0;32m/home/spike/ml/udacity/nd101/deep-learning-modified/tv-script-generation/problem_unittests.py\u001b[0m in \u001b[0;36mtest_pick_word\u001b[0;34m(pick_word)\u001b[0m\n\u001b[1;32m    283\u001b[0m         \u001b[0mtest_int_to_vocab\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0mword_i\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mword\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mword_i\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mword\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'this'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'is'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'a'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'test'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    284\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 285\u001b[0;31m         \u001b[0mpred_word\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpick_word\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtest_probabilities\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtest_int_to_vocab\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    286\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    287\u001b[0m         \u001b[0;31m# Check type\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m<ipython-input-472-ec56654dfaa2>\u001b[0m in \u001b[0;36mpick_word\u001b[0;34m(probabilities, int_to_vocab)\u001b[0m\n\u001b[1;32m     10\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchoice\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprobabilities\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mp\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mprobabilities\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 12\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mint_to_vocab\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mprobabilities\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchoice\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprobabilities\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mp\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mprobabilities\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     13\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;31mIndexError\u001b[0m: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices"
      ]
     }
    ],
@@ -2029,9 +2094,10 @@
     "    :return: String of the predicted word\n",
     "    \"\"\"\n",
     "    \n",
-    "    word = int_to_vocab[np.argmax(probabilities)]\n",
+    "    word_p = np.random.choice(probabilities,p=probabilities)\n",
+    "    word = np.extract(probabilities==word_p, probabilities)\n",
     "    \n",
-    "    return word\n",
+    "    return int_to_vocab[probabilities[np.random.choice(probabilities,p=probabilities)]]\n",
     "\n",
     "\n",
     "\"\"\"\n",