# Buggy code
def forward_diag(self, x1, x2):
    lengthscale = self.log_lengthscale.exp()
    mean = x1.mean(1).mean(0)
    x1_normed = (x1 - mean.unsqueeze(0).unsqueeze(1)).div(lengthscale)
    x2_normed = (x2 - mean.unsqueeze(0).unsqueeze(1)).div(lengthscale)

    diff = x1_normed - x2_normed
    distance_over_rho = diff.pow_(2).sum(-1).sqrt()
    exp_component = torch.exp(-math.sqrt(self.nu * 2) * distance_over_rho)

    if self.nu == 0.5:
        constant_component = 1
    elif self.nu == 1.5:
        constant_component = (math.sqrt(3) * distance_over_rho).add(1)
    elif self.nu == 2.5:
        constant_component = (
            (math.sqrt(5) * distance_over_rho)
            .add(1)
            .add(5.0 / 3.0 * distance_over_rho**2)
        )

    return constant_component * exp_component


# Traceback (RuntimeError)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-144-12b763efd692> in <module>()
19     output = model(train_x)
20     # TODO: Fix this view call!!
---> 21     loss = -mll(output, train_y)
22     loss.backward()
23     print('Iter %d/%d - Loss: %.3f' % (i + 1, n_iter, loss.item()))

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/module.py in __call__(self, *inputs, **kwargs)
160
161     def __call__(self, *inputs, **kwargs):
--> 162         outputs = self.forward(*inputs, **kwargs)
163         if torch.is_tensor(outputs) or isinstance(outputs, RandomVariable) or isinstance(outputs, LazyVariable):
164             return outputs

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/mlls/exact_marginal_log_likelihood.py in forward(self, output, target)
49
50         # Get log determininat and first part of quadratic form
---> 51         inv_quad, log_det = covar.inv_quad_log_det(inv_quad_rhs=(target - mean).unsqueeze(-1), log_det=True)
52
53         # Add terms for SGPR / when inducing points are learned

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/lazy/lazy_variable.py in inv_quad_log_det(self, inv_quad_rhs, log_det)
577         matrix_size = self.size(-1)
578         batch_size = self.size(0) if self.ndimension() == 3 else None
--> 579         tensor_cls = self.tensor_cls
580
581         args = lazy_var.representation()

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/lazy/lazy_variable.py in tensor_cls(self)
882     def tensor_cls(self):
883         if not hasattr(self, "_tensor_cls"):
--> 884             first_item = self.representation()[0]
885             if isinstance(first_item, Variable):
886                 first_item = first_item.data

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/lazy/lazy_variable.py in representation(self)
733                 representation.append(arg)
734             elif isinstance(arg, LazyVariable):
--> 735                 representation += list(arg.representation())
736             else:
737                 raise RuntimeError("Representation of a LazyVariable should consist only of Variables")

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/lazy/lazy_variable.py in representation(self)
733                 representation.append(arg)
734             elif isinstance(arg, LazyVariable):
--> 735                 representation += list(arg.representation())
736             else:
737                 raise RuntimeError("Representation of a LazyVariable should consist only of Variables")

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/lazy/lazy_evaluated_kernel_variable.py in representation(self)
109
110     def representation(self):
--> 111         return self.evaluate_kernel().representation()
112
113     def representation_tree(self):

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/lazy/lazy_evaluated_kernel_variable.py in evaluate_kernel(self)
96                 x2 = self.x2
97
---> 98             self._cached_kernel_eval = super(Kernel, self.kernel).__call__(x1, x2, **self.params)
99             if self.squeeze_row:
100                 self._cached_kernel_eval.squeeze_(-2)

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/module.py in __call__(self, *inputs, **kwargs)
160
161     def __call__(self, *inputs, **kwargs):
--> 162         outputs = self.forward(*inputs, **kwargs)
163         if torch.is_tensor(outputs) or isinstance(outputs, RandomVariable) or isinstance(outputs, LazyVariable):
164             return outputs

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/kernels/multitask_kernel.py in forward(self, x1, x2)
61     def forward(self, x1, x2):
62         covar_i = self.task_covar_module.covar_matrix
---> 63         covar_x = self.data_covar_module.forward(x1, x2)
64         if covar_x.size(0) == 1:
65             covar_x = covar_x[0]

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/kernels/grid_interpolation_kernel.py in forward(self, x1, x2, **kwargs)
52
53     def forward(self, x1, x2, **kwargs):
---> 54         base_lazy_var = self._inducing_forward()
55         if x1.size(0) > 1:
56             base_lazy_var = base_lazy_var.repeat(x1.size(0), 1, 1)

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/kernels/grid_interpolation_kernel.py in _inducing_forward(self)
46     def _inducing_forward(self):
47         inducing_points_var = Variable(self.inducing_points)
---> 48         return super(GridInterpolationKernel, self).forward(inducing_points_var, inducing_points_var)
49
50     def forward_diag(self, x1, x2, **kwargs):

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/kernels/grid_kernel.py in forward(self, x1, x2, **kwargs)
38             if settings.use_toeplitz.on():
39                 first_item = grid_var[:, 0:1].contiguous()
---> 40                 covar_columns = self.base_kernel_module(first_item, grid_var, **kwargs).evaluate()
41                 covars = [ToeplitzLazyVariable(covar_columns[i : i + 1].squeeze(-2)) for i in range(n_dim)]
42             else:

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/lazy/lazy_evaluated_kernel_variable.py in evaluate(self)
115
116     def evaluate(self):
--> 117         return self.evaluate_kernel().evaluate()
118
119     def exact_predictive_mean(self, full_mean, train_labels, n_train, likelihood, precomputed_cache=None):

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/lazy/lazy_evaluated_kernel_variable.py in evaluate_kernel(self)
96                 x2 = self.x2
97
---> 98             self._cached_kernel_eval = super(Kernel, self.kernel).__call__(x1, x2, **self.params)
99             if self.squeeze_row:
100                 self._cached_kernel_eval.squeeze_(-2)

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/module.py in __call__(self, *inputs, **kwargs)
160
161     def __call__(self, *inputs, **kwargs):
--> 162         outputs = self.forward(*inputs, **kwargs)
163         if torch.is_tensor(outputs) or isinstance(outputs, RandomVariable) or isinstance(outputs, LazyVariable):
164             return outputs

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/kernels/rbf_kernel.py in forward(self, x1, x2)
104     def forward(self, x1, x2):
105         lengthscales = self.log_lengthscale.exp().mul(math.sqrt(2)).clamp(self.eps, 1e5)
--> 106         diff = (x1.unsqueeze(2) - x2.unsqueeze(1)).div_(lengthscales.unsqueeze(1))
107         return diff.pow_(2).sum(-1).mul_(-1).exp_()

RuntimeError: The expanded size of the tensor (1) must match the existing size (2) at non-singleton dimension 3

# Fixed code
def forward_diag(self, x1, x2):
    mean = x1.mean(1, keepdim=True).mean(0, keepdim=True)
    x1_normed = x1 - mean.unsqueeze(0).unsqueeze(1)
    x2_normed = x2 - mean.unsqueeze(0).unsqueeze(1)

    diff = x1_normed - x2_normed
    distance_over_rho = diff.pow_(2).sum(-1).sqrt()
    exp_component = torch.exp(-math.sqrt(self.nu * 2) * distance_over_rho)

    if self.nu == 0.5:
        constant_component = 1
    elif self.nu == 1.5:
        constant_component = (math.sqrt(3) * distance_over_rho).add(1)
    elif self.nu == 2.5:
        constant_component = (
            (math.sqrt(5) * distance_over_rho)
            .add(1)
            .add(5.0 / 3.0 * distance_over_rho**2)
        )

    return constant_component * exp_component