diff --git a/Sequence Models/Neural machine translation with attention - v2.ipynb b/Sequence Models/Neural machine translation with attention - v2.ipynb
index 57c9c7e4..35736f89 100644
--- a/Sequence Models/Neural machine translation with attention - v2.ipynb	
+++ b/Sequence Models/Neural machine translation with attention - v2.ipynb	
@@ -866,7 +866,7 @@
     "for example in EXAMPLES:\n",
     "    \n",
     "    source = string_to_int(example, Tx, human_vocab)\n",
-    "    source = np.array(list(map(lambda x: to_categorical(x, num_classes=len(human_vocab)), source))).swapaxes(0,1)\n",
+    "    source = np.array(list(map(lambda x: to_categorical(x, num_classes=len(human_vocab)), source)), ndmin=3)\n",
     "    prediction = model.predict([source, s0, c0])\n",
     "    prediction = np.argmax(prediction, axis = -1)\n",
     "    output = [inv_machine_vocab[int(i)] for i in prediction]\n",