diff --git a/.github/workflows/package_test.yml b/.github/workflows/package_test.yml
index af8c7506..b99469db 100644
--- a/.github/workflows/package_test.yml
+++ b/.github/workflows/package_test.yml
@@ -94,32 +94,15 @@ jobs:
 
     - name: Install the package and run pytest
       timeout-minutes: 10
-      shell: bash -l {0}
       run: |
-        # install the wheel package
-        # sudo apt install python3-venv
-
-        # extract cuda version
-        CUDA_VERSION=$(nvcc --version | grep "Cuda compilation tools" | cut -d ' ' -f 5 | cut -d ',' -f 1)
-
-        # Create and activate a virtual environment
-        python3 -m venv .venv
-        source .venv/bin/activate
-
-        # Install PyTorch
-        pip install torch==${TORCH_VERSION} -i "https://download.pytorch.org/whl/cu${CUDA_VERSION//./}"
-
-        # Install the wheel package
-        pip install dist/*.whl
-
-        # Install the test requirements
-        pip install -r requirements-test.txt
-        
-        # Run pytest
-        printf "\n\nRunning pytest\n\n"
-        cd tests
-        python3 -m pytest || exit 1
-        printf "\n\n"
-
-        # Deactivate the virtual environment
-        deactivate
+        docker pull pytorch/manylinux-builder:cuda${CUDA_VERSION}
+        docker run --rm -t --gpus=all \
+            -v "$CI_CACHE_DIR":/ci_cache \
+            -v "$GITHUB_WORKSPACE":/ScaleLLM \
+            -e PYTHON_VERSION=${PYTHON_VERSION} \
+            -e CUDA_VERSION=${CUDA_VERSION} \
+            -e TORCH_VERSION=${TORCH_VERSION} \
+            -e PIP_CACHE_DIR=/ci_cache/.pip \
+            -u $(id -u):$(id -g) \
+            pytorch/manylinux-builder:cuda${CUDA_VERSION} \
+            bash /ScaleLLM/scripts/run_pytest.sh
diff --git a/scripts/run_pytest.sh b/scripts/run_pytest.sh
new file mode 100755
index 00000000..8011b202
--- /dev/null
+++ b/scripts/run_pytest.sh
@@ -0,0 +1,41 @@
+#!/bin/bash
+set -e
+
+ensure_env() {
+    local var_name="$1"
+    if [ -z "${!var_name}" ]; then
+        echo "Error: Environment variable '$var_name' is not set."
+        exit 1
+    fi
+}
+
+ensure_env PYTHON_VERSION
+ensure_env TORCH_VERSION
+ensure_env CUDA_VERSION
+
+PROJECT_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
+
+export HOME=/tmp/home
+mkdir -p $HOME
+export PATH="$HOME/.local/bin:$PATH"
+
+# choose python version
+PYVER="${PYTHON_VERSION//./}"
+export PATH="/opt/python/cp${PYVER}-cp${PYVER}/bin:$PATH"
+
+# install PyTorch
+pip install torch==$TORCH_VERSION -i "https://download.pytorch.org/whl/cu${CUDA_VERSION//./}"
+
+cd "$PROJECT_ROOT"
+
+# install dependencies
+pip install numpy
+pip install -r requirements-test.txt
+
+# install scalellm wheel
+pip install dist/*.whl
+
+# run pytest
+printf "\n\nRunning pytest\n\n"
+# cd tests
+python3 -m pytest
\ No newline at end of file