Weird behavior for manual independent Multi-Output GP

[Turakar]

I try to model a two-output problem using a single ExactGP with two kernel functions. As such, the resulting kernel matrix is the block diagonal matrix of the two individual kernel matrices. This is the corresponding code:

class IndependentModel(gpytorch.models.ExactGP):
    """
    Class models a multi-output setting with multiple separate kernels with no cross-correlation.
    As such, the kernel matrix is a block diagonal matrix.
    """
    
    def __init__(self, train_x, train_y, likelihood, num_tasks):
        super().__init__(train_x, train_y, likelihood)
        self.mean_modules = ModuleList([
            gpytorch.means.ZeroMean() for _ in range(num_tasks)
        ])
        self.covar_modules = ModuleList([
            gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel()) for _ in range(num_tasks)
        ])

    def forward(self, x, i):
        # Order samples by ascending task index
        i = i.squeeze()  # remove the last dimension added by GPyTorch
        order = torch.argsort(i, stable=True)
        x_ordered = x[order]
        unorder = torch.arange(len(x))[order]  # create a reverse mapping
        
        # Calculate offsets of task indices in the data
        num_tasks = len(self.mean_modules)
        task_offsets = torch.cumsum(torch.tensor(
            [torch.tensor(0)] + [torch.count_nonzero(torch.eq(i, task)) for task in range(num_tasks)]
        ), 0)
        
        # Calculate and concatenate means
        mean_x = torch.cat([
            self.mean_modules[task](x_ordered[task_offsets[task] : task_offsets[task + 1]]).to_dense() for task in range(num_tasks)
        ])[unorder]
        
        # Calculate and concatenate covariances
        covar_x = torch.block_diag(*[
            self.covar_modules[task](x_ordered[task_offsets[task] : task_offsets[task + 1]]).to_dense() for task in range(num_tasks)
        ])[unorder, :][:, unorder]
        
        # Return final distribution
        return MultivariateNormal(mean_x, linear_operators.to_linear_operator(covar_x))

See here for a full example on a synthetic dataset. Note that the Plotly plots only render if you open the notebook in JupyterLab and not on GitHub. As you can see there, this implementation does not work. If I simultaneously ask for predictions for both tasks, I get different results than if I ask for predictions per task.

I think this is a numerical stability problem caused by the first process not properly converging (cf. plots at the end). However, I cannot find the actual problem that causes this. Maybe one of you can help me with finding the root cause here?

I am aware that I could use the model list implemented in GPyTorch, but as I want to move on to correlated outputs, I first want to get this simple example working.