Fix device mismatch issue in pytorch#1071 (pytorch#1073)

* fix device mismatch issue pytorch#1071

Author: hudeven

imagenet/main.py

@@ -7,18 +7,18 @@
 from enum import Enum
 
 import torch
-import torch.nn as nn
-import torch.nn.parallel
 import torch.backends.cudnn as cudnn
 import torch.distributed as dist
-import torch.optim
-from torch.optim.lr_scheduler import StepLR
 import torch.multiprocessing as mp
+import torch.nn as nn
+import torch.nn.parallel
+import torch.optim
 import torch.utils.data
 import torch.utils.data.distributed
-import torchvision.transforms as transforms
 import torchvision.datasets as datasets
 import torchvision.models as models
+import torchvision.transforms as transforms
+from torch.optim.lr_scheduler import StepLR
 from torch.utils.data import Subset
 
 model_names = sorted(name for name in models.__dict__
@@ -275,13 +275,12 @@ def main_worker(gpu, ngpus_per_node, args):
             train_sampler.set_epoch(epoch)
 
         # train for one epoch
-        train(train_loader, model, criterion, optimizer, epoch, args)
+        train(train_loader, model, criterion, optimizer, epoch, device, args)
 
         # evaluate on validation set
         acc1 = validate(val_loader, model, criterion, args)
 
         scheduler.step()
-
 
         # remember best acc@1 and save checkpoint
         is_best = acc1 > best_acc1
@@ -299,7 +298,7 @@ def main_worker(gpu, ngpus_per_node, args):
             }, is_best)
 
 
-def train(train_loader, model, criterion, optimizer, epoch, args):
+def train(train_loader, model, criterion, optimizer, epoch, device, args):
     batch_time = AverageMeter('Time', ':6.3f')
     data_time = AverageMeter('Data', ':6.3f')
     losses = AverageMeter('Loss', ':.4e')
@@ -318,13 +317,9 @@ def train(train_loader, model, criterion, optimizer, epoch, args):
         # measure data loading time
         data_time.update(time.time() - end)
 
-        if args.gpu is not None and torch.cuda.is_available():
-            images = images.cuda(args.gpu, non_blocking=True)
-        elif not args.gpu and torch.cuda.is_available():
-            target = target.cuda(args.gpu, non_blocking=True)
-        elif torch.backends.mps.is_available():
-            images = images.to('mps')
-            target = target.to('mps')
+        # move data to the same device as model
+        images = images.to(device, non_blocking=True)
+        target = target.to(device, non_blocking=True)
 
         # compute output
         output = model(images)