diff --git a/tensorflow_privacy/privacy/membership_inference_attack/codelab.ipynb b/tensorflow_privacy/privacy/membership_inference_attack/codelab.ipynb
index b809e69..e717432 100644
--- a/tensorflow_privacy/privacy/membership_inference_attack/codelab.ipynb
+++ b/tensorflow_privacy/privacy/membership_inference_attack/codelab.ipynb
@@ -101,7 +101,6 @@
         "  pass\n",
         "\n",
         "import numpy as np\n",
-        "from colabtools import adhoc_import\n",
         "from typing import Tuple, Text\n",
         "from scipy import special\n",
         "\n",
@@ -132,7 +131,7 @@
       },
       "outputs": [],
       "source": [
-        "!pip install tensorflow_privacy\n",
+        "!pip3 install git+https://github.com/tensorflow/privacy\n",
         "\n",
         "from tensorflow_privacy.privacy.membership_inference_attack import membership_inference_attack as mia"
       ]
@@ -252,8 +251,8 @@
         "cce = tf.keras.backend.categorical_crossentropy\n",
         "constant = tf.keras.backend.constant\n",
         "\n",
-        "loss_train = cce(constant(y_train), prob_train, from_logits=False).numpy()\n",
-        "loss_test = cce(constant(y_test), prob_test, from_logits=False).numpy()"
+        "loss_train = cce(constant(y_train), constant(prob_train), from_logits=False).numpy()\n",
+        "loss_test = cce(constant(y_test), constant(prob_test), from_logits=False).numpy()"
       ]
     },
     {