diff --git a/docker/dockerfile.ctr b/docker/dockerfile.ctr
index f79b5ed85..ea459669a 100644
--- a/docker/dockerfile.ctr
+++ b/docker/dockerfile.ctr
@@ -77,6 +77,7 @@ RUN if [[ "${HUGECTR_DEV_MODE}" == "false" ]]; then \
         mv /hugectr/ci ~/hugectr-ci && rm -rf /hugectr && mkdir -p /hugectr && mv ~/hugectr-ci /hugectr/ci \
     ; fi
 
+ENV PYTHONPATH=${PYTHONPATH}:${HUGECTR_HOME}/lib
 ARG _HUGECTR_BACKEND_REPO="github.com/triton-inference-server/hugectr_backend.git"
 ARG TRITON_VERSION
 # Install Triton inference backend.
diff --git a/docker/dockerfile.tf b/docker/dockerfile.tf
index d69576691..e86129268 100644
--- a/docker/dockerfile.tf
+++ b/docker/dockerfile.tf
@@ -17,8 +17,8 @@ COPY --chown=1000:1000 --from=triton /opt/tritonserver/backends/tensorflow backe
 # Tensorflow dependencies (only)
 # Pinning to pass hugectr sok tests
 # wrapt 1.5.0 introduce hugectr test failures, so downgrade to 1.14.0
-RUN pip install --no-cache-dir tensorflow protobuf==3.20.3 wrapt==1.14.0 \
-    && pip uninstall tensorflow keras -y
+RUN pip install --no-cache-dir tensorflow==2.12.0 protobuf==3.20.3 wrapt==1.14.0 \
+     && pip uninstall tensorflow keras -y
 
 # DLFW Tensorflow packages
 COPY --chown=1000:1000 --from=dlfw /usr/local/lib/python${PYTHON_VERSION}/dist-packages/tensorflow /usr/local/lib/python${PYTHON_VERSION}/dist-packages/tensorflow/