added missing cuda call in pytorch model (#24)

bethgelab · Jun 20, 2017 · cb27d1c · cb27d1c
1 parent 02230ed
commit cb27d1c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/foolbox/models/pytorch.py b/foolbox/models/pytorch.py
@@ -65,6 +65,8 @@ def predictions_and_gradient(self, image, label):
 
         target = np.array([label])
         target = torch.from_numpy(target)
+        if self.cuda:  # pragma: no cover
+            target = target.cuda()
         target = Variable(target)
 
         assert image.ndim == 3