From 1bee58f5bbf7d7b2364a8cf338fb9757e202385b Mon Sep 17 00:00:00 2001
From: alex-mccarthy-unity
 <122620276+alex-mccarthy-unity@users.noreply.github.com>
Date: Wed, 6 Mar 2024 14:58:26 +0100
Subject: [PATCH] Fix cuda not being used on windows: update pytorch version
 (#6064)

These references were missed when upgrading from pytorch 1.x to 2.x in #6013

References found by running `grep -R '1\.13\.1' .`

Install command chosen from the guide at https://pytorch.org/get-started/locally/
---
 .yamato/pytest-gpu.yml | 2 +-
 docs/Installation.md   | 6 ++++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/.yamato/pytest-gpu.yml b/.yamato/pytest-gpu.yml
index 248d33079..2ab36e718 100644
--- 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.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
+      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
       if python -c "exec('import torch \nif not torch.cuda.is_available(): raise')" &> /dev/null; then
         echo 'all good'
       else
diff --git a/docs/Installation.md b/docs/Installation.md
index 8f240ef13..73efad9ab 100644
--- a/docs/Installation.md
+++ b/docs/Installation.md
@@ -146,10 +146,12 @@ offer a dedicated [guide on Virtual Environments](Using-Virtual-Environment.md).
 #### (Windows) Installing PyTorch
 
 On Windows, you'll have to install the PyTorch package separately prior to
-installing ML-Agents. Activate your virtual environment and run from the command line:
+installing ML-Agents in order to make sure the cuda-enabled version is used,
+rather than the CPU-only version. Activate your virtual environment and run from
+the command line:
 
 ```sh
-pip3 install torch~=1.13.1 -f https://download.pytorch.org/whl/torch_stable.html
+pip3 install torch~=2.2.1 --index-url https://download.pytorch.org/whl/cu121
 ```
 
 Note that on Windows, you may also need Microsoft's