Fix issue with .diag() when broadcasting SumLazyTensors
Looks like we were trying to be too smart here, this would cause RuntimeErrors when having base lazy tensors of different dimension.
Instead we can simply rely on torch tensor broadcasting here.
@@ -77,8 +77,4 @@ class SumLazyTensor(LazyTensor):
raise AttributeError("other must be a LazyTensor")
def diag(self):
- diags = [lazy_tensor.diag().contiguous() for lazy_tensor in self.lazy_tensors]
- size = diags[0].size()
- res = sum(diag.view(-1) for diag in diags)
- res = res.view(size)
- return res
+ return sum(lazy_tensor.diag().contiguous() for lazy_tensor in self.lazy_tensors)
