Replace cuda.device with env var to assign GPU to processes

This commit is contained in:
Yin Li 2020-07-22 18:48:12 -07:00
parent 98cdd4795c
commit ca39b2028f

View File

@ -42,8 +42,12 @@ def node_worker(args):
def gpu_worker(local_rank, node, args):
device = torch.device('cuda', local_rank)
torch.cuda.device(device)
#device = torch.device('cuda', local_rank)
#torch.cuda.device(device) # env var recommended over this
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = str(local_rank)
device = torch.device('cuda', 0)
rank = args.gpus_per_node * node + local_rank