diff --git a/hierarchical_probabilistic_unet/HPU_Net.ipynb b/hierarchical_probabilistic_unet/HPU_Net.ipynb
index 4352b07..fcf3e1b 100644
--- a/hierarchical_probabilistic_unet/HPU_Net.ipynb
+++ b/hierarchical_probabilistic_unet/HPU_Net.ipynb
@@ -39,8 +39,7 @@
     "id": "ntSZbVLBTzIC"
    },
    "source": [
-    "#Colab: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/googlecolab/colabtools/blob/master/notebooks/colab-github-demo.ipynb)\n",
-    "\n",
+    "#Colab: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deepmind/deepmind-research/blob/master/hierarchical_probabilistic_unet/HPU_Net.ipynb)\n",
     "If you haven't already opened this notebook in colab, you can click the button above to open it there."
    ]
   },
@@ -59,6 +58,7 @@
     "!pip install dm-sonnet==1.35\n",
     "\n",
     "from glob import glob\n",
+    "import matplotlib\n",
     "import matplotlib.pyplot as plt\n",
     "import matplotlib.gridspec as gridspec\n",
     "import numpy as np\n",
@@ -87,9 +87,9 @@
    },
    "outputs": [],
    "source": [
-    "!git clone https://github.com/deepmind/deepmind-research.git .\n",
-    "!pip install ./hierarchical_prob_unet\n",
-    "from model import HierarchicalProbUNet"
+    "!git clone https://github.com/deepmind/deepmind-research.git deepmind_research\n",
+    "%cd deepmind_research/hierarchical_probabilistic_unet\n",
+    "from model import HierarchicalProbUNet\n"
    ]
   },
   {
@@ -112,8 +112,8 @@
    },
    "outputs": [],
    "source": [
-    "!gsutil cp gs://hpunet_data/test/test.tar.gz /tmp/\n",
-    "!tar xvfz /tmp/test.tar.gz -C /tmp/test_data"
+    "!gsutil cp gs://hpunet-data/lidc_crops/test.tar.gz /tmp/\n",
+    "!tar xvfz /tmp/test.tar.gz -C /tmp/"
    ]
   },
   {
@@ -136,8 +136,8 @@
    },
    "outputs": [],
    "source": [
-    "!gsutil cp gs://hpunet_data/model_weights/checkpoint.tar.gz /tmp/\n",
-    "!tar xvfz /tmp/checkpoint.tar.gz -C /tmp/checkpoint"
+    "!gsutil cp gs://hpunet-data/model_checkpoint/checkpoint.tar.gz /tmp/\n",
+    "!tar xvfz /tmp/checkpoint.tar.gz -C /tmp/"
    ]
   },
   {
@@ -153,7 +153,7 @@
    "source": [
     "#@title Utility to function to batch-up the data.\n",
     "\n",
-    "test_img_dir = '/tmp/test_data'\n",
+    "test_img_dir = '/tmp/test'\n",
     "\n",
     "def make_batch(patient_ixs):\n",
     "  \"\"\"Assembles a batch of the first image for each specified patient.\n",
@@ -164,7 +164,7 @@
     "    segmentations of shape (b, 128, 128, 4).\n",
     "  \"\"\"\n",
     "  img, seg = [], []\n",
-    "  patient_dirs = glob(os.path.join(test_img_dir, '*'))\n",
+    "  patient_dirs = glob(os.path.join(test_img_dir, 'images', '*'))\n",
     "  for i in patient_ixs:\n",
     "    # get the first image for the patient\n",
     "    img_path = glob(os.path.join(patient_dirs[i], '*'))[0]\n",
@@ -254,7 +254,7 @@
     "saver = tf.train.Saver()\n",
     "with tf.train.MonitoredTrainingSession() as sess:\n",
     "  # restore checkpoint\n",
-    "  saver.restore(sess, '/tmp/checkpoint')\n",
+    "  saver.restore(sess, '/tmp/checkpoint/tf_graph_data.ckpt')\n",
     "\n",
     "  # sample from the full hierarchy\n",
     "  rec, s = sess.run([reconstruction, sample_full],\n",
diff --git a/hierarchical_probabilistic_unet/README.md b/hierarchical_probabilistic_unet/README.md
index f8f543d..6270a0a 100644
--- a/hierarchical_probabilistic_unet/README.md
+++ b/hierarchical_probabilistic_unet/README.md
@@ -40,8 +40,8 @@ link to the colab.
 ## Installation
 To install the package locally run:
 ```bash
-git clone https://github.com/deepmind/hierarchical_prob_unet .
-cd hierarchical_prob_unet
+git clone https://github.com/deepmind/deepmind-research.git .
+cd deepmind-research/hierarchical_probabilistic_unet
 pip install -e .
 ```