Fix ConstantDiagLT creation
@@ -28,7 +28,7 @@ class LowRankRootAddedDiagLazyTensor(AddedDiagLazyTensor):
A_inv = self._diag_tensor.inverse() # This is fine since it's a DiagLazyTensor
U = self._lazy_tensor.root
V = self._lazy_tensor.root.transpose(-2, -1)
- C = ConstantDiagLazyTensor(torch.ones(V.batch_shape, device=V.device, dtype=V.dtype), V.shape[-2])
+ C = ConstantDiagLazyTensor(torch.ones(*V.batch_shape, 1, device=V.device, dtype=V.dtype), V.shape[-2])
cap_mat = delazify(C + V.matmul(A_inv.matmul(U)))
chol_cap_mat = psd_safe_cholesky(cap_mat)
