diff --git a/.github/workflows/gpu.yml b/.github/workflows/gpu.yml
index 0a7c65435f..8c6b2c515e 100644
--- a/.github/workflows/gpu.yml
+++ b/.github/workflows/gpu.yml
@@ -27,10 +27,6 @@ jobs:
       - uses: actions/checkout@v3
         with:
           fetch-depth: 0
-      - name: Install Ubuntu packages
-        run: |
-          apt-get update -y
-          apt-get install -y lsb-release
       - name: Install and upgrade python packages
         run: |
           python -m pip install --upgrade pip tox
@@ -54,15 +50,17 @@ jobs:
   gpu-cu11:
     runs-on: linux-amd64-gpu-p100-latest-1
     container:
-      image: nvidia/cuda:11.8.0-devel-ubuntu22.04
+      image: nvidia/cuda:11.8.0-runtime-ubuntu22.04
       env:
         NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
     steps:
       - name: Install Ubuntu packages
         run: |
           apt-get update -y
-          # libcudnn8 installed for tensorflow GPU support
-          apt-get install -y git lsb-release 'libcudnn8=*cuda11.8'
+          apt-get install -y \
+            git \
+            'libcudnn8=*cuda11.8' `# tensorflow GPU support` \
+            cuda-nvcc-11-8 `# required for numba`
       - uses: actions/checkout@v3
         with:
           fetch-depth: 0
@@ -83,7 +81,7 @@ jobs:
         uses: NVIDIA-Merlin/.github/actions/branch-name@6f0539fba24f60da2aee63c5925bee7cee3206e3
       - name: Run tests
         run: |
-          if [[ "${{ github.ref }}" != 'refs/heads/main' ]]; then
+          if [ "${{ github.ref }}" != 'refs/heads/main' ]; then
               extra_pytest_markers="and changed"
           fi
           merlin_branch="${{ steps.get-branch-name.outputs.branch }}"
@@ -94,7 +92,7 @@ jobs:
   tests-examples:
     runs-on: linux-amd64-gpu-p100-latest-1
     container:
-      image: nvidia/cuda:11.8.0-devel-ubuntu22.04
+      image: nvidia/cuda:11.8.0-runtime-ubuntu22.04
       env:
         NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
     steps:
@@ -102,7 +100,10 @@ jobs:
         run: |
           apt-get update -y
           # libcudnn8 installed for tensorflow GPU support
-          apt-get install -y git lsb-release 'libcudnn8=*cuda11.8'
+          apt-get install -y \
+            git \
+            'libcudnn8=*cuda11.8' `# tensorflow GPU support` \
+            cuda-nvcc-11-8 `# required for numba`
       - uses: actions/checkout@v3
         with:
           fetch-depth: 0