How to do Hadamard Multitask GP Regression with batched input data

[EricLangezaal]

Hi! I have been trying to do Hadamard Multitask GP Regression, for batched input data. Copying the tutorial exactly allows me to correctly perform the regression on a single vector input x.

I cannot get it to work for batched inputs of the shape [batch, num_datapoints, dim_x]. I have taken a look at this tutorial, but still couldn't get it to work.

My current code is the following:

likelihood = gpytorch.likelihoods.GaussianLikelihood(batch_shape=xc.shape[:1])

xc_labels = torch.concat((
    torch.zeros(10, 1),
    torch.ones(10, 1), 
    dim=0
).to(torch.long)

model = MultitaskGPModel((xc, xc_labels), yc, self.kernel, likelihood, out_dim=2)
model.eval()

outputDist = likelihood(model(xt, torch.zeros(xt.shape[-2], 1 , dtype=torch.long)))

class MultitaskGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, kernel, likelihood, out_dim):
        super(MultitaskGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ZeroMean(batch_shape=train_x[0].shape[:1])
        self.covar_module = kernel
        self.covar_module.batch_shape = train_x[0].shape[:1]
        self.task_module = gpytorch.kernels.IndexKernel(num_tasks=out_dim)

    def forward(self, x, i):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        covar_i = self.task_module(i)
        covar = covar_x.mul(covar_i)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar)

Where kernel is a simple RBF kernel and xc is of shape [batch, num_data, x_dim] and yc of shape [batch,num_data, 1]. This fails with the error:

Flattening the training labels failed. The most common cause of this error is that the shapes of the prior mean and the training labels are mismatched. The shape of the train targets is torch.Size([1, 110, 1]), while the reported shape of the mean is torch.Size([1, 110])

If I instead use an xc of shape [num_data] (i.e. no batches and 1-dimensional x_dim) and do not set the batch_shape anywhere it works just fine, but I really want to be able to do this for batched inputs.

Does anyone know how to properly do batched Hadamard Multitask GP regression?