Fix for GPU runtime error for semiring (#124)

urchade · web-flow · commit 7b3b661063d2 · 2022-01-30T13:56:48.000-05:00
Fix #121
diff --git a/torch_struct/semirings/semirings.py b/torch_struct/semirings/semirings.py
@@ -49,6 +49,7 @@ def dot(cls, a, b):
 
     @staticmethod
     def fill(c, mask, v):
+        mask = mask.to(c.device)
         return torch.where(
             mask, v.type_as(c).view((-1,) + (1,) * (len(c.shape) - 1)), c
         )

Original file line number	Diff line number	Diff line change
`@@ -49,6 +49,7 @@ def dot(cls, a, b):`
`49`	`49`
`50`	`50`	`@staticmethod`
`51`	`51`	`def fill(c, mask, v):`
	`52`	`+ mask = mask.to(c.device)`
`52`	`53`	`return torch.where(`
`53`	`54`	`mask, v.type_as(c).view((-1,) + (1,) * (len(c.shape) - 1)), c`
`54`	`55`	`)`