diff --git a/test/test_artifacts/v0/keras.test.Dockerfile b/test/test_artifacts/v0/keras.test.Dockerfile
index 74690920..2843f14d 100644
--- a/test/test_artifacts/v0/keras.test.Dockerfile
+++ b/test/test_artifacts/v0/keras.test.Dockerfile
@@ -8,7 +8,7 @@ RUN sudo apt-get update && sudo apt-get install -y git graphviz && \
     :
 
 # Some of the keras guides requires pydot and graphviz to be installed
-RUN micromamba install -y --freeze-installed conda-forge::pydot nvidia::cuda-nvcc
+RUN micromamba install -y --freeze-installed -c conda-forge pydot cuda-nvcc jax
 ENV XLA_FLAGS=--xla_gpu_cuda_data_dir=/opt/conda
 
 WORKDIR "keras-io/guides"
diff --git a/test/test_artifacts/v0/run_pandas_tests.py b/test/test_artifacts/v0/run_pandas_tests.py
index 08953750..971066a6 100644
--- a/test/test_artifacts/v0/run_pandas_tests.py
+++ b/test/test_artifacts/v0/run_pandas_tests.py
@@ -3,6 +3,7 @@
 import sys
 
 import pandas
+import boto3
 
 # We change the working directory here because there is at least one test (`test_html_template_extends_options`) which
 # expects the directory to be 'pandas'. Ideally, we would have changed directories through a `WORKDIR` in Dockerfile
diff --git a/test/test_artifacts/v0/scripts/run_pysdk_tests.sh b/test/test_artifacts/v0/scripts/run_pysdk_tests.sh
index 2f49f122..a90d51bd 100644
--- a/test/test_artifacts/v0/scripts/run_pysdk_tests.sh
+++ b/test/test_artifacts/v0/scripts/run_pysdk_tests.sh
@@ -2,6 +2,9 @@
 
 # We need to checkout the version of sagemaker-python-sdk that is installed in the mamba environment.
 
+# Sagemaker Python SDK's unit tests requires AWS_DEFAULT_REGION to be set. So, using an arbitrary value of us-east-1
+export AWS_DEFAULT_REGION=us-west-2
+
 pysdk_version=$(micromamba list | grep sagemaker-python-sdk | tr -s ' ' | cut -d ' ' -f 3)
 # Checkout the corresponding sagemaker-python-sdk version
 git checkout tags/v$pysdk_version
diff --git a/test/test_artifacts/v0/sm-python-sdk.test.Dockerfile b/test/test_artifacts/v0/sm-python-sdk.test.Dockerfile
index 1e9426d1..ad0942fd 100644
--- a/test/test_artifacts/v0/sm-python-sdk.test.Dockerfile
+++ b/test/test_artifacts/v0/sm-python-sdk.test.Dockerfile
@@ -5,8 +5,7 @@ ARG MAMBA_DOCKERFILE_ACTIVATE=1
 RUN sudo apt-get update && sudo apt-get install -y git && \
     git clone --recursive https://github.com/aws/sagemaker-python-sdk.git && \
     :
-# Sagemaker Python SDK's unit tests requires AWS_DEFAULT_REGION to be set. So, using an arbitrary value of us-east-1
-ENV AWS_DEFAULT_REGION=us-east-1
+
 WORKDIR "sagemaker-python-sdk"
 COPY --chown=$MAMBA_USER:$MAMBA_USER scripts/run_pysdk_tests.sh .
 RUN chmod +x run_pysdk_tests.sh
diff --git a/test/test_artifacts/v1/keras.test.Dockerfile b/test/test_artifacts/v1/keras.test.Dockerfile
index f5cb224f..3f07277d 100644
--- a/test/test_artifacts/v1/keras.test.Dockerfile
+++ b/test/test_artifacts/v1/keras.test.Dockerfile
@@ -8,7 +8,7 @@ RUN sudo apt-get update && sudo apt-get install -y git graphviz && \
     :
 
 # Some of the keras guides requires pydot and graphviz to be installed
-RUN micromamba install -y --freeze-installed conda-forge::pydot "nvidia::cuda-nvcc>=11.8,<11.9"
+RUN micromamba install -y --freeze-installed conda-forge::pydot "nvidia::cuda-nvcc>=11.8,<11.9" jax
 ENV XLA_FLAGS=--xla_gpu_cuda_data_dir=/opt/conda
 
 WORKDIR "keras-io/guides"
diff --git a/test/test_artifacts/v1/scripts/run_pysdk_tests.sh b/test/test_artifacts/v1/scripts/run_pysdk_tests.sh
index 2f49f122..a90d51bd 100644
--- a/test/test_artifacts/v1/scripts/run_pysdk_tests.sh
+++ b/test/test_artifacts/v1/scripts/run_pysdk_tests.sh
@@ -2,6 +2,9 @@
 
 # We need to checkout the version of sagemaker-python-sdk that is installed in the mamba environment.
 
+# Sagemaker Python SDK's unit tests requires AWS_DEFAULT_REGION to be set. So, using an arbitrary value of us-east-1
+export AWS_DEFAULT_REGION=us-west-2
+
 pysdk_version=$(micromamba list | grep sagemaker-python-sdk | tr -s ' ' | cut -d ' ' -f 3)
 # Checkout the corresponding sagemaker-python-sdk version
 git checkout tags/v$pysdk_version
diff --git a/test/test_artifacts/v1/sm-python-sdk.test.Dockerfile b/test/test_artifacts/v1/sm-python-sdk.test.Dockerfile
index 623efbf6..aa628ab3 100644
--- a/test/test_artifacts/v1/sm-python-sdk.test.Dockerfile
+++ b/test/test_artifacts/v1/sm-python-sdk.test.Dockerfile
@@ -4,8 +4,6 @@ FROM $SAGEMAKER_DISTRIBUTION_IMAGE
 ARG MAMBA_DOCKERFILE_ACTIVATE=1
 RUN git clone --recursive https://github.com/aws/sagemaker-python-sdk.git
 
-# Sagemaker Python SDK's unit tests requires AWS_DEFAULT_REGION to be set. So, using an arbitrary value of us-east-1
-ENV AWS_DEFAULT_REGION=us-east-1
 WORKDIR "sagemaker-python-sdk"
 COPY --chown=$MAMBA_USER:$MAMBA_USER scripts/run_pysdk_tests.sh .
 RUN chmod +x run_pysdk_tests.sh