fix n_gpu count when no_cuda flag is activated (#3077)

* fix n_gpu count when no_cuda flag is activated

* someone was left behind
This commit is contained in:
Victor SANH
2020-03-02 10:20:21 -05:00
committed by GitHub
parent 298bed16a8
commit 6b1ff25084
13 changed files with 13 additions and 13 deletions

View File

@@ -535,7 +535,7 @@ def main():
# Setup CUDA, GPU & distributed training
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
args.n_gpu = torch.cuda.device_count()
args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()
else: # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)