diff --git a/.yamato/pytest-gpu.yml b/.yamato/pytest-gpu.yml
index 2ab36e7182..dd7d3161f0 100644
--- a/.yamato/pytest-gpu.yml
+++ b/.yamato/pytest-gpu.yml
@@ -11,10 +11,11 @@ pytest_gpu:
       python3 -m pip install pyyaml --index-url https://artifactory.prd.it.unity3d.com/artifactory/api/pypi/pypi/simple
       python3 -u -m ml-agents.tests.yamato.setup_venv
       python3 -m pip install --progress-bar=off -r test_requirements.txt --index-url https://artifactory.prd.it.unity3d.com/artifactory/api/pypi/pypi/simple
-      python3 -m pip install torch==2.2.1+cu121 torchvision==0.17.1+cu121 torchaudio==0.17.1 --index-url https://download.pytorch.org/whl/cu121
+      python3 -m pip install torch==2.2.1+cu118 torchvision==0.17.1+cu118 torchaudio==2.2.1+cu118 --index-url https://download.pytorch.org/whl/cu118
       if python -c "exec('import torch \nif not torch.cuda.is_available(): raise')" &> /dev/null; then
         echo 'all good'
       else
+        echo 'cuda device not available!'
         exit 1
       fi
       python3 -m pytest -m "not slow" --junitxml=junit/test-results.xml -p no:warnings