Update main.py

Normalize train and validate if statements

Author: Quentin-Anthony

imagenet/main.py

@@ -290,11 +290,11 @@ 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
         output = model(images)
         loss = criterion(output, target)
@@ -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