Переглянути джерело

Updating async to non_blocking in cuda calls, for Pytorch 0.4+ compatibility and python >= 3.7

Jason Antic 6 роки тому
батько
коміт
871111c69a
2 змінених файлів з 2 додано та 2 видалено
  1. 1 1
      fastai/core.py
  2. 1 1
      fastai/models/cifar10/main_dxy.py

+ 1 - 1
fastai/core.py

@@ -34,7 +34,7 @@ def T(a, half=False, cuda=True):
         elif a.dtype in (np.float32, np.float64):
             a = torch.cuda.HalfTensor(a) if half else torch.FloatTensor(a)
         else: raise NotImplementedError(a.dtype)
-    if cuda: a = to_gpu(a, async=True)
+    if cuda: a = to_gpu(a, non_blocking=True)
     return a
 
 def create_variable(x, volatile, requires_grad=False):

+ 1 - 1
fastai/models/cifar10/main_dxy.py

@@ -220,7 +220,7 @@ def validate(val_loader, model, criterion, log):
 
   for i, (input, target) in enumerate(val_loader):
     if args.use_cuda:
-      target = target.cuda(async=True)
+      target = target.cuda(non_blocking=True)
       input = input.cuda()
     input_var = torch.autograd.Variable(input, volatile=True)
     target_var = torch.autograd.Variable(target, volatile=True)