diff --git a/demo/notebooks/tree_inspection.ipynb b/demo/notebooks/tree_inspection.ipynb
index 0c9149c4..089baf49 100644
--- a/demo/notebooks/tree_inspection.ipynb
+++ b/demo/notebooks/tree_inspection.ipynb
@@ -134,7 +134,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "forest_preds_y_mcmc = bart_model.y_hat_test[:,bart_model.num_gfr:]\n",
+    "forest_preds_y_mcmc = bart_model.y_hat_test\n",
     "y_avg_mcmc = np.squeeze(forest_preds_y_mcmc).mean(axis = 1, keepdims = True)\n",
     "y_df_mcmc = pd.DataFrame(np.concatenate((np.expand_dims(y_test,1), y_avg_mcmc), axis = 1), columns=[\"True outcome\", \"Average estimated outcome\"])\n",
     "sns.scatterplot(data=y_df_mcmc, x=\"Average estimated outcome\", y=\"True outcome\")\n",