[Bug] Multitask GP Regression fails for LeaveOneOutPseudoLikelihood mll

[sahakpetros]

🐛 Bug

To reproduce

run Multitask GP Regression example but replace ExactMarginalLogLikelihood with LeaveOneOutPseudoLikelihood
** Code snippet to reproduce **
in Multitask GP Regression replace

mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

with

mll = gpytorch.mlls.LeaveOneOutPseudoLikelihood(likelihood, model)

** Stack trace/error message **

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[5], line 20
     18 optimizer.zero_grad()
     19 output = model(train_x)
---> 20 loss = -mll(output, train_y)
     21 loss.backward()
     22 print('Iter %d/%d - Loss: %.3f' % (i + 1, training_iterations, loss.item()))

File ~/.venv/lib/python3.11/site-packages/gpytorch/module.py:31, in Module.__call__(self, *inputs, **kwargs)
     30 def __call__(self, *inputs, **kwargs) -> Union[Tensor, Distribution, LinearOperator]:
---> 31     outputs = self.forward(*inputs, **kwargs)
     32     if isinstance(outputs, list):
     33         return [_validate_module_outputs(output) for output in outputs]

File ~/.venv/lib/python3.11/site-packages/gpytorch/mlls/leave_one_out_pseudo_likelihood.py:62, in LeaveOneOutPseudoLikelihood.forward(self, function_dist, target, *params)
     60 identity = torch.eye(*L.shape[-2:], dtype=m.dtype, device=m.device)
     61 sigma2 = 1.0 / L._cholesky_solve(identity, upper=False).diagonal(dim1=-1, dim2=-2)  # 1 / diag(inv(K))
---> 62 mu = target - L._cholesky_solve((target - m).unsqueeze(-1), upper=False).squeeze(-1) * sigma2
     63 term1 = -0.5 * sigma2.log()
     64 term2 = -0.5 * (target - mu).pow(2.0) / sigma2

File ~/.venv/lib/python3.11/site-packages/linear_operator/operators/triangular_linear_operator.py:79, in TriangularLinearOperator._cholesky_solve(self, rhs, upper)
     72 def _cholesky_solve(
     73     self: Float[LinearOperator, "*batch N N"],
...
     36     upper: Optional[bool] = False,
     37 ) -> Union[Float[LinearOperator, "... N M"], Float[Tensor, "... N M"]]:
---> 38     return torch.cholesky_solve(rhs, self.to_dense(), upper=upper)

RuntimeError: Incompatible matrix sizes for cholesky_solve: each A matrix is 200 by 200 but each b matrix is 2 by 1

Expected Behavior

no RuntimeError

System information

Please complete the following information:

• GPyTorch 1.13
• PyTorch 2.4.1+cpu
• ChromeOS 128

Additional context

this bug can be fixed by flattening "target" variable