Normalize train and validate GPU movement conditions

Author: Quentin Anthony

imagenet/main.py

@@ -290,9 +290,9 @@ def train(train_loader, model, criterion, optimizer, epoch, args):
         # measure data loading time
         data_time.update(time.time() - end)
 
-        if torch.cuda.is_available():
+        if args.gpu is not None:
             images = images.cuda(args.gpu, non_blocking=True)
-        if torch.cuda.is_available():
+        if args.gpu is not None:
             target = target.cuda(args.gpu, non_blocking=True)
 
         # compute output
@@ -334,9 +334,9 @@ def validate(val_loader, model, criterion, args):
     with torch.no_grad():
         end = time.time()
         for i, (images, target) in enumerate(val_loader):
-            if torch.cuda.is_available():
+            if args.gpu is not None:
                 images = images.cuda(args.gpu, non_blocking=True)
-            if torch.cuda.is_available():
+            if args.gpu is not None:
                 target = target.cuda(args.gpu, non_blocking=True)
 
             # compute output