Update Docker image and run commands in workflow

Author: Jyothirmaikottu

.github/workflows/test-cdk-changes.yaml

@@ -16,9 +16,8 @@ jobs:
       - name: Check Docker
         run: |
           docker info
-          echo "Pod NVIDIA_VISIBLE_DEVICES: $NVIDIA_VISIBLE_DEVICES"
-          docker pull nvidia/cuda:12.0.0-base-ubuntu22.04
-          docker run --rm --gpus "\"device=${NVIDIA_VISIBLE_DEVICES}\"" nvidia/cuda:12.0.0-base-ubuntu22.04 nvidia-smi
+          docker pull nvidia/cuda:12.0.0-runtime-ubuntu22.04
+          docker run --rm --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all nvidia/cuda:12.0.0-runtime-ubuntu22.04 nvidia-smi
 
   test-p4d:
     runs-on: p4d-gpu-runners
@@ -29,6 +28,5 @@ jobs:
       - name: Check Docker
         run: |
           docker info
-          echo "Pod NVIDIA_VISIBLE_DEVICES: $NVIDIA_VISIBLE_DEVICES"
-          docker pull nvidia/cuda:12.0.0-base-ubuntu22.04
-          docker run --rm --gpus "\"device=${NVIDIA_VISIBLE_DEVICES}\"" nvidia/cuda:12.0.0-base-ubuntu22.04 nvidia-smi
+          docker pull nvidia/cuda:12.0.0-runtime-ubuntu22.04
+          docker run --rm --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all nvidia/cuda:12.0.0-runtime-ubuntu22.04 nvidia-smi