Updated yamato pytest GPU tests for torch 1.13.1 (#5985)

Unity-Technologies · Oct 6, 2023 · 63b71a8 · 63b71a8
1 parent fb24597
commit 63b71a8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/.yamato/pytest-gpu.yml b/.yamato/pytest-gpu.yml
@@ -11,7 +11,7 @@ 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==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html --index-url https://artifactory.prd.it.unity3d.com/artifactory/api/pypi/pypi/simple
+      python3 -m pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu117 --index-url https://artifactory.prd.it.unity3d.com/artifactory/api/pypi/pypi/simple
       if python -c "exec('import torch \nif not torch.cuda.is_available(): raise')" &> /dev/null; then
         echo 'all good'
       else