diff --git a/Week 3 Programming Assignment.ipynb b/Week 3 Programming Assignment.ipynb
index 61bda69..9247853 100644
--- a/Week 3 Programming Assignment.ipynb	
+++ b/Week 3 Programming Assignment.ipynb	
@@ -845,12 +845,15 @@
     "    # I couldn't get this one right. Let me know if you figured it out.\n",
     "    if initial_state is None:\n",
     "        model.layers[1].reset_states()\n",
+    "    #Below code is right\n"
     "    else:\n",
-    "        initial_state = model.layers[1].states\n",
+    "        model.layers[1].reset_states(states=initial_state)\n",
     "        \n",
-    "    prediction = model.predict(token_sequence)\n",
+    "    token_sequence = np.array(token_sequence)\n"
+    "    predicted = model.predict(token_sequence)\n"
+    "    pred = np.expand_dims(predicted[0,-1,:],0)\n",
     "    \n",
-    "    return prediction[0]"
+    "    return pred"
    ]
   },
   {
@@ -916,10 +919,8 @@
     "    distribution and return the sample as a single integer.\n",
     "    \"\"\"\n",
     "    # I couldn't get this one right. Let me know if you figured it out.\n",
-    "    sample = tf.random.categorical(logits,1)\n",
-    "    sampe = tf.squeeze(sample,axis = -1)\n",
-    "    \n",
-    "    return sample[0]"
+    "    # Below code is right
+    "    return int(tf.random.categorical(logits,1))"
    ]
   },
   {