using LeaveOneOutPseudoLikelihood mll for Multitask GP Regression

[partev]

When I run the Multitask GP Regression example from here:
https://docs.gpytorch.ai/en/stable/examples/03_Multitask_Exact_GPs/Multitask_GP_Regression.html

and replace ExactMarginalLogLikelihood with LeaveOneOutPseudoLikelihood

mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)
replaced with
mll = gpytorch.mlls.LeaveOneOutPseudoLikelihood(likelihood, model)

I get the following error:

{
	"name": "RuntimeError",
	"message": "Incompatible matrix sizes for cholesky_solve: each A matrix is 200 by 200 but each b matrix is 2 by 1",
	"stack": "---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[6], 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:78, in TriangularLinearOperator._cholesky_solve(self, rhs, upper)
     71 def _cholesky_solve(
     72     self: Float[LinearOperator, \"*batch N N\"],
     73     rhs: Union[Float[LinearOperator, \"*batch2 N M\"], Float[Tensor, \"*batch2 N M\"]],
     74     upper: Optional[bool] = False,
     75 ) -> Union[Float[LinearOperator, \"... N M\"], Float[Tensor, \"... N M\"]]:
     76     # use custom method if implemented
     77     try:
---> 78         res = self._tensor._cholesky_solve(rhs=rhs, upper=upper)
     79     except NotImplementedError:
     80         if upper:
     81             # res = (U.T @ U)^-1 @ v = U^-1 @ U^-T @ v

File ~/.venv/lib/python3.11/site-packages/linear_operator/operators/dense_linear_operator.py:38, in DenseLinearOperator._cholesky_solve(self, rhs, upper)
     33 def _cholesky_solve(
     34     self: Float[LinearOperator, \"*batch N N\"],
     35     rhs: Union[Float[LinearOperator, \"*batch2 N M\"], Float[Tensor, \"*batch2 N M\"]],
     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"
}

Is LeaveOneOutPseudoLikelihood mll fundamentally incompatible with Multitask GPR or is this a bug that can be fixed?