Multitask kernel with ApproximateGP

[flcello]

I want to use a Multitask kernel as in the Multitask GP Regression Tutorial, but with ApproximateGPs. However, I struggle to figure out which wrappers and batch shapes to use exactly. Could someone help me out?

import torch
import gpytorch

class MultitaskGPModel(gpytorch.models.ApproximateGP):
    def __init__(self, num_tasks, num_inducing_points=10):

        inducing_points = torch.rand(num_tasks, num_inducing_points, 1)

        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
            num_inducing_points,
            batch_shape=torch.Size([num_tasks]),
        )

        variational_strategy = gpytorch.variational.IndependentMultitaskVariationalStrategy(
            gpytorch.variational.VariationalStrategy(
                self,
                inducing_points,
                variational_distribution,
            ),
            num_tasks,
            task_dim=-1,
        )

        super().__init__(variational_strategy)

        self.mean_module = gpytorch.means.MultitaskMean(gpytorch.means.ConstantMean(), num_tasks)
        self.covar_module = gpytorch.kernels.MultitaskKernel(gpytorch.kernels.RBFKernel(), num_tasks)

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultitaskMultivariateNormal(mean_x, covar_x)

gpmodel = MultitaskGPModel(num_tasks=2)

x = torch.randn(20, 1)
gpmodel(x).sample()