Add pass to cpu

deep-spin · Sep 8, 2021 · dd01298 · dd01298
1 parent ed96c93
commit dd01298
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/lpsmap/api/autograd.py b/lpsmap/api/autograd.py
@@ -17,8 +17,8 @@ class _LPSparseMAP(torch.autograd.Function):
 
     @classmethod
     def forward(cls, ctx, fg, eta_u, *eta_v):
-        fg.set_log_potentials(eta_u.detach().numpy())
-        detached_eta_v = [np.atleast_1d(x.detach().numpy()) for x in eta_v]
+        fg.set_log_potentials(eta_u.detach().cpu().numpy())
+        detached_eta_v = [np.atleast_1d(x.detach().cpu().numpy()) for x in eta_v]
         fg.set_all_additionals(detached_eta_v)
         ctx.fg = fg
         ctx.shape = eta_u.shape
@@ -30,7 +30,7 @@ def backward(cls, ctx, du):
         dtype = du.dtype
         device = du.device
 
-        du = du.to(dtype=torch.double, device="cpu").detach().numpy()
+        du = du.to(dtype=torch.double, device="cpu").detach().cpu().numpy()
         out = torch.empty(ctx.shape, dtype=torch.double, device='cpu')
         add = ctx.fg.jacobian_vec(du, out.numpy())