diff --git a/advanced_functionality/pytorch_deploy_large_GPT_model/GPT-J-6B-model-parallel-inference-DJL.ipynb b/advanced_functionality/pytorch_deploy_large_GPT_model/GPT-J-6B-model-parallel-inference-DJL.ipynb
index 74565437de..da11c73cb7 100644
--- a/advanced_functionality/pytorch_deploy_large_GPT_model/GPT-J-6B-model-parallel-inference-DJL.ipynb
+++ b/advanced_functionality/pytorch_deploy_large_GPT_model/GPT-J-6B-model-parallel-inference-DJL.ipynb
@@ -114,8 +114,7 @@
     "\n",
     "\n",
     "# # Build the docker image locally with the image name and then push it to ECR\n",
-    "image_id=$(docker images -q | head -n1)\n",
-    "docker tag $image_id ${fullname}\n",
+    "docker tag deepjavalibrary/djl-serving:0.18.0-deepspeed ${fullname}\n",
     "\n",
     "docker push $fullname"
    ]