Minor train and README updates for install on Mac and code location

README.md

@@ -11,6 +11,8 @@ This code currently only replicates the claw environment experiments of the pape
 - Install Python (ran on Python 3.9) and/or create a fresh environment
 - Install requirements `pip install -r requirements.txt`
 
+For Mac you would need to remove the requirement for CUDA from `requirements.txt` for `torchvision` and install it with CPU support.
+
 ## Running
 
 ### Mini script
@@ -22,6 +24,7 @@ As a lightweight entry-point to the code, we provide `claw_mini_script.py`, whic
 To recreate the claw machine figures in the paper, run `./run.sh` or alternatively run each step separately:
 
 ```
+cd code
 python make_dataset.py
 python train.py
 python plot.py
@@ -40,4 +43,4 @@ Use of Microsoft trademarks or logos in modified versions of this project must n
 Any use of third-party trademarks or logos are subject to those third-party's policies.
 
 ## License
-Code is licensed under MIT, data and all other content is licensed under Microsoft Research License Agreement (MSR-LA). See LICENSE.
+Code is licensed under MIT, data and all other content is licensed under Microsoft Research License Agreement (MSR-LA). See LICENSE.

code/train.py

@@ -83,7 +83,7 @@
 
 n_epoch = 100
 lrate = 1e-4
-device = "cuda"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 n_hidden = 512
 batch_size = 32
 n_T = 50