Fix batch GPs
@@ -100,6 +100,7 @@ class ExactGP(Module):
n_tasks = 1
if isinstance(full_output, MultitaskGaussianRandomVariable):
n_tasks = full_output.n_tasks
+
if self.train_targets.ndimension() == 2:
# Multitask
n_train = self.train_targets.size(0)
