fix grad on cuda

Author: unixpickle

reptile_gen/batching.py

@@ -18,7 +18,7 @@ def run_grad_fn(inputs, outputs):
         if device != 'cpu':
             model.to(d)
         res = grad_fn(model, inputs.to(d), outputs.to(d))
-        return [p.grad for p in model.parameters()], res
+        return [p.grad.cpu() for p in model.parameters()], res
 
     pool = Pool(min(len(batch), threads))
     pickled_fn = cloudpickle.dumps(run_grad_fn)