diff --git a/tools/colab/shakespeare_with_tpu_and_keras.ipynb b/tools/colab/shakespeare_with_tpu_and_keras.ipynb
index 8f35dd5f..36b83ba5 100644
--- a/tools/colab/shakespeare_with_tpu_and_keras.ipynb
+++ b/tools/colab/shakespeare_with_tpu_and_keras.ipynb
@@ -423,7 +423,7 @@
         "  \n",
         "  # sample from our output distribution\n",
         "  next_idx = [\n",
-        "      np.random.choice(256, p=next_probits[i])\n",
+        "      [np.random.choice(256, p=next_probits[i])]\n",
         "      for i in range(BATCH_SIZE)\n",
         "  ]\n",
         "  predictions.append(np.asarray(next_idx, dtype=np.int32))\n",
@@ -432,7 +432,7 @@
         "for i in range(BATCH_SIZE):\n",
         "  print('PREDICTION %d\\n\\n' % i)\n",
         "  p = [predictions[j][i] for j in range(PREDICT_LEN)]\n",
-        "  generated = ''.join([chr(c) for c in p])  # Convert back to text\n",
+        "  generated = ''.join([chr(c[0]) for c in p])  # Convert back to text\n",
         "  print(generated)\n",
         "  print()\n",
         "  assert len(generated) == PREDICT_LEN, 'Generated text too short'"