diff --git a/docker/tritonbench-nightly.dockerfile b/docker/tritonbench-nightly.dockerfile
index 148df0cc..aa312276 100644
--- a/docker/tritonbench-nightly.dockerfile
+++ b/docker/tritonbench-nightly.dockerfile
@@ -22,13 +22,13 @@ RUN cd /workspace/tritonbench && \
 
 RUN cd /workspace/tritonbench && \
     . ${SETUP_SCRIPT} && \
-    sudo python tools/cuda_utils.py --setup-cuda-softlink
+    sudo python -m tools.cuda_utils --setup-cuda-softlink
 
 # Install PyTorch nightly and verify the date is correct
 RUN cd /workspace/tritonbench && \
     . ${SETUP_SCRIPT} && \
-    python tools/cuda_utils.py --install-torch-deps && \
-    python tools/cuda_utils.py --install-torch-nightly
+    python -m tools.cuda_utils --install-torch-deps && \
+    python -m tools.cuda_utils --install-torch-nightly
 
 # Check the installed version of nightly if needed
 RUN cd /workspace/tritonbench && \
@@ -37,9 +37,9 @@ RUN cd /workspace/tritonbench && \
         echo "torch version check skipped"; \
     elif [ -z "${FORCE_DATE}" ]; then \
         FORCE_DATE=$(date '+%Y%m%d') \
-        python tools/cuda_utils.py --check-torch-nightly-version --force-date "${FORCE_DATE}"; \
+        python -m tools.cuda_utils --check-torch-nightly-version --force-date "${FORCE_DATE}"; \
     else \
-        python tools/cuda_utils.py --check-torch-nightly-version --force-date "${FORCE_DATE}"; \
+        python -m tools.cuda_utils --check-torch-nightly-version --force-date "${FORCE_DATE}"; \
     fi
 
 # Tritonbench library build and test require libcuda.so.1
diff --git a/docker/tritonbench-rocm-nightly.dockerfile b/docker/tritonbench-rocm-nightly.dockerfile
index 21bbfde7..5fb5f3e6 100644
--- a/docker/tritonbench-rocm-nightly.dockerfile
+++ b/docker/tritonbench-rocm-nightly.dockerfile
@@ -35,8 +35,8 @@ RUN cd /workspace/tritonbench && \
 # Install PyTorch nightly and verify the date is correct
 RUN cd /workspace/tritonbench && \
     . ${SETUP_SCRIPT} && \
-    python tools/rocm_utils.py --install-torch-deps && \
-    python tools/rocm_utils.py --install-torch-nightly
+    python -m tools.rocm_utils --install-torch-deps && \
+    python -m tools.rocm_utils --install-torch-nightly
 
 
 # Install Tritonbench