Fix forward_diag
@@ -39,10 +39,10 @@ class MultitaskKernel(Kernel):
base data kernel and the task kernel.
"""
task_indices = torch.arange(self.n_tasks, device=x1.device).long()
- task_indices = task_indices.unsqueeze(-1).unsqueeze(-1)
+ task_indices = task_indices.unsqueeze(0).unsqueeze(-1)
# These are small because they are vectors, therefore it is safe to evaluate them
- covar_i_diag = self.task_covar_module.forward_diag(task_indices)
+ covar_i_diag = self.task_covar_module.forward_diag(task_indices, task_indices)
covar_x_diag = self.data_covar_module.forward_diag(x1, x2)
if isinstance(covar_x_diag, LazyVariable):
