def METHOD_NAME(self, x, i):
    # Get predictive mean
    mean_x = self.mean_module(x)
    # Get all covariances, we'll look up the task-speicific ones
    covar_x = self.covar_module(x)
    # # Get the covariance for task i
    covar_i = self.task_covar_module(i)
    covar_xi = covar_x.mul(covar_i)
    return MultivariateNormal(mean_x, covar_xi)