Allow user to specify the list of GPUs to run on using CUDA_VISIBLE_DEVICES environment variable.
By default, the current code gets the list of GPUs using torch.cuda.device_count() and runs all detected GPUs
However, say if we want to use GPU 2 and 3 only on a 4-GPU machine with the command: CUDA_VISIBLE_DEVICES=2,3 python3 ..., the code still uses GPU 0 and 1 because it specifies the GPU for each command using something like ['CUDA_VISIBLE_DEVICES={str(x)}' for x in range(torch.cuda.device_count())]
Thus, I add the code to check for the CUDA_VISIBLE_DEVICES environment variable and specify the corresponding GPU IDs in the command_launcher.py file. In this way, GPUs 2 and 3 are used instead of 0 and 1.
Allow user to specify the list of GPUs to run on using
CUDA_VISIBLE_DEVICES
environment variable.torch.cuda.device_count()
and runs all detected GPUsCUDA_VISIBLE_DEVICES=2,3 python3 ...
, the code still uses GPU 0 and 1 because it specifies the GPU for each command using something like['CUDA_VISIBLE_DEVICES={str(x)}' for x in range(torch.cuda.device_count())]
CUDA_VISIBLE_DEVICES
environment variable and specify the corresponding GPU IDs in thecommand_launcher.py
file. In this way, GPUs 2 and 3 are used instead of 0 and 1.Let me know if I can help more!