Simplify chol lazy tensor checks
@@ -12,10 +12,7 @@ class CholLazyTensor(RootLazyTensor):
chol = chol.evaluate()
# Check that we have a lower triangular matrix
- mask = torch.full((chol.size(-2), chol.size(-2)), fill_value=-1, dtype=chol.dtype, device=chol.device)
- mask.tril_().add_(1)
- if chol.ndimension() == 3:
- mask.unsqueeze_(0)
+ mask = torch.ones(chol.shape[-2:], dtype=chol.dtype, device=chol.device).triu_(1)
if torch.max(chol.mul(mask)).item() > 1e-3 and torch.equal(chol, chol):
raise RuntimeError("CholLazyVaraiable should take a lower-triangular " "matrix in the constructor.")
