update readme

GETTING_STARTED.md

@@ -68,12 +68,19 @@ To set up a virtual enviornment and install this repository
 
     For **PyTorch**
 
+    Note: the below command assumes you have CUDA 12.1 installed locally.
+    This is the default in the provided Docker image.
+    We recommend you match this CUDA version but if you decide to run
+    with a different local CUDA version, please find the appropriate wheel
+    url to pass to the `pip install` command for `pytorch`.
+
     ```bash
     pip3 install -e '.[jax_cpu]'
-    pip3 install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/torch_stable.html'
+    pip3 install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/cu121'
     pip3 install -e '.[full]'
     ```
 
+
 <details>
 <summary>
 Per workload installations

setup.cfg

@@ -137,9 +137,11 @@ pytorch_cpu =
   torchvision==0.16.0
 
 # PyTorch GPU
+# Note: omit the cuda suffix and installing from the appropriate
+# wheel will result in using locally installed CUDA.
 pytorch_gpu =
-  torch==2.1.0+cu118
-  torchvision==0.16.0+cu118
+  torch==2.1.0
+  torchvision==0.16.0
 
 # wandb
 wandb =