.

torch_struct/semirings.py

@@ -140,8 +140,9 @@ def unaccumulate_(a, b, grad_output, fn, step=10000):
     # b_grad = b.clone().fill_(0)
     # print("chcek", a_grad.shape)    
     a_grad = torch.tensor(0.0, device=a.device).set_(a.clone().storage(), a.storage_offset(), a.size(), a.stride()).fill_(0)
-
+    a_grad = a_grad.expand_as(a)
     b_grad = torch.tensor(0.0, device=b.device).set_(b.clone().storage(), b.storage_offset(), b.size(), b.stride()).fill_(0)
+    b_grad = b_grad.expand_as(b)
     print(b_grad.shape, a.shape)
 
     print("chcek", a_grad.shape)