diff --git a/.github/workflows/test_3d_parallelism.yaml b/.github/workflows/test_3d_parallelism.yaml
index fefddbc5..96a52e2b 100644
--- a/.github/workflows/test_3d_parallelism.yaml
+++ b/.github/workflows/test_3d_parallelism.yaml
@@ -20,16 +20,17 @@ jobs:
   tests:
     runs-on: [multi-gpu, nvidia-gpu, 8-t4, ci]
     container:
-      image: nvcr.io/nvidia/pytorch:23.03-py3
+      # image: nvcr.io/nvidia/pytorch:23.03-py3
+      image: runpod/pytorch:2.1.1-py3.10-cuda12.1.1-devel-ubuntu22.04
       ports:
         - 80
       options: --gpus all --shm-size "8G"
     steps:
     - uses: actions/checkout@v3
-    - name: Set up Python 3.10
-      uses: actions/setup-python@v2
-      with:
-        python-version: '3.10'
+    # - name: Set up Python 3.10
+    #   uses: actions/setup-python@v2
+    #   with:
+    #     python-version: '3.10'
 
     - name: Python environment
       run: |