.

harvardnlp · Oct 31, 2019 · 6909467 · 6909467
1 parent 331eb7f
commit 6909467
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torch_struct/semirings.py b/torch_struct/semirings.py
@@ -139,9 +139,9 @@ def unaccumulate_(a, b, grad_output, fn, step=10000):
     # a_grad = a.clone().fill_(0)    
     # b_grad = b.clone().fill_(0)
     # print("chcek", a_grad.shape)    
-    a_grad = torch.tensor(0.0).set_(a.clone().storage(), a.storage_offset(), a.size(), a.stride()).fill_(0)
+    a_grad = torch.tensor(0.0, device=a.device).set_(a.clone().storage(), a.storage_offset(), a.size(), a.stride()).fill_(0)
 
-    b_grad = torch.tensor(0.0).set_(b.clone().storage(), b.storage_offset(), b.size(), b.stride()).fill_(0)
+    b_grad = torch.tensor(0.0, device=b.device).set_(b.clone().storage(), b.storage_offset(), b.size(), b.stride()).fill_(0)
 
     print("chcek", a_grad.shape)
     total = 1