Updated arg documentation, added --cpu to train just in case

README.md

@@ -42,6 +42,17 @@ Each folder contains 8 replicates with varying performance.
 
 ### Included scripts
 
+See the docstring for each script for a full list of arguments (debug, other testing args).
+
+Common (technical) arguments:
+
+| Arg | Type | Description |
+| - | - | - |
+| ```gpus``` | ```int[]``` | Comma separated list of GPUs (1) |
+| ```cpu``` | ```bool``` | Whether to run on CPU instead of GPU |
+
+(1) GPUs are specified by GPU index (i.e. as returned by ```gpustat```). If no ```--gpus``` are provided, all GPUs on the system are used. If no GPUs are installed, CPU will be used.
+
 ```evaluate.py```:
 
 | Arg | Type | Description |
@@ -50,8 +61,6 @@ Each folder contains 8 replicates with varying performance.
 | ```directory``` | ```str``` | Target directory to load from. Can pass a comma separated list. |
 | ```repeat``` | ```int``` | Number of times to run evaluation. Default: 10 |
 
-See the help string for a full list of arguments.
-
 ```train.py```: 
 
 | Arg | Type | Description |
@@ -68,8 +77,6 @@ See the help string for a full list of arguments.
 | ```problem``` | ```str``` | Problem to evaluate on. Can pass a comma separated list. |
 | ```optimizer``` | ```str``` | Name of optimizer to use. |
 
-See the help string for a full list of arguments.
-
 ### Experiment folder structure
 
 Experiment file path:

train.py

@@ -11,6 +11,8 @@
     are created, and a mirrored strategy is created with all physical GPUs.
 --vram : int
     Amount of vram to allocate per virtual GPU if testing.
+--cpu : bool
+    Whether to run on CPU instead of GPU.
 --gpus : int[]
     Comma separated list of GPU indices to use on a multi-gpu system.
 --initialize : bool
@@ -45,8 +47,9 @@
 vgpus = int(args.pop_get("--vgpu", default=1))
 memory_limit = int(args.pop_get("--vram", default=12000))
 gpus = args.pop_get("--gpus", default=None)
+cpu = args.pop_get("--cpu", default=False, dtype=bool)
 distribute = create_distribute(
-    vgpus=vgpus, memory_limit=memory_limit, gpus=gpus)
+    vgpus=vgpus, memory_limit=memory_limit, gpus=gpus, do_cpu=cpu)
 
 # Pick up flags first
 initialize_only = args.pop_check("--initialize")