fix cuda

Author: gabrielhuang

models.py

@@ -13,16 +13,18 @@ def point_grad_to(self, target):
         Set .grad attribute of each parameter to be proportional
         to the difference between self and target
         '''
-        is_cuda = next(self.parameters()).is_cuda
         for p, target_p in zip(self.parameters(), target.parameters()):
             if p.grad is None:
-                if is_cuda:
+                if self.is_cuda():
                     p.grad = Variable(torch.zeros(p.size())).cuda()
                 else:
                     p.grad = Variable(torch.zeros(p.size()))
             p.grad.data.zero_()  # not sure this is required
             p.grad.data.add_(p.data - target_p.data)
 
+    def is_cuda(self):
+        return next(self.parameters()).is_cuda
+
 
 class OmniglotModel(ReptileModel):
     """
@@ -77,6 +79,8 @@ def predict(self, prob):
     def clone(self):
         clone = OmniglotModel(self.num_classes)
         clone.load_state_dict(self.state_dict())
+        if self.is_cuda():
+            clone.cuda()
         return clone