diff --git a/build_scripts/inference/serve b/build_scripts/inference/serve
index 962abe52..6efa32f9 100755
Binary files a/build_scripts/inference/serve and b/build_scripts/inference/serve differ
diff --git a/middleware_api/trainings/create_training_job.py b/middleware_api/trainings/create_training_job.py
index a7f21f2f..9e92adae 100644
--- a/middleware_api/trainings/create_training_job.py
+++ b/middleware_api/trainings/create_training_job.py
@@ -137,6 +137,10 @@ def _trigger_sagemaker_training_job(
             "s3_location": f"s3://{bucket_name}/{train_params_file}",
         },
         job_id=train_job.id,
+        environment={
+            "SITE_PACKAGES_S3_PATH": f"aws-gcr-solutions-{region}/"
+                                     f"stable-diffusion-aws-extension-github-mainline/{esd_version}/train.tar"
+        }
     )
     est.fit(wait=False)