diff --git a/advanced_functionality/causal-inference/causal-inference-container.ipynb b/advanced_functionality/causal-inference/causal-inference-container.ipynb
index 7dbd3a41fe..208549e10a 100644
--- a/advanced_functionality/causal-inference/causal-inference-container.ipynb
+++ b/advanced_functionality/causal-inference/causal-inference-container.ipynb
@@ -108,7 +108,7 @@
     "fi\n",
     "\n",
     "# Get the login command from ECR and execute it directly\n",
-    "aws ecr get-login-password --region ${region}|docker login --username AWS --password-stdin ${fullname}\n",
+    "$(aws ecr get-login --region $region --registry-ids $account --no-include-email)\n",
     "\n",
     "# Build the docker image locally with the image name and then push it to ECR\n",
     "# with the full name.\n",
@@ -176,6 +176,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "! mkdir data\n",
+    "\n",
     "# S3 bucket where the training data is located.\n",
     "data_bucket = f\"sagemaker-sample-files\"\n",
     "data_prefix = \"datasets/tabular/uci_heart_failure/\"\n",