Aquila-Consortium/map2map
32
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Yin Li 2020-07-22 18:48:12 -07:00
parent 98cdd4795c
commit ca39b2028f
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
map2map/train.py
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