# Buggy code
def forward(self, x1, x2):
    batch_size, n, n_dims = x1.size()
    _, m, _ = x2.size()
    if not n_dims == self.n_dims:
        raise RuntimeError("The number of dimensions doesn't match what was supplied!")

    mixture_weights = self.log_mixture_weights.view(self.n_mixtures, 1, 1, 1).exp()
    mixture_means = self.log_mixture_means.view(
        self.n_mixtures, 1, 1, 1, self.n_dims
    ).exp()
    mixture_scales = self.log_mixture_scales.view(
        self.n_mixtures, 1, 1, 1, self.n_dims
    ).exp()
    distance = (x1.unsqueeze(-2) - x2.unsqueeze(-3)).abs()  # distance = x^(i) - z^(i)

    exp_term = (distance * mixture_scales).pow_(2).mul_(-2 * math.pi**2)
    cos_term = (distance * mixture_means).mul_(2 * math.pi)
    res = exp_term.exp_() * cos_term.cos_()

    # Product over dimensions
    res = res.prod(-1)

    # Sum over mixtures
    res = (res * mixture_weights).sum(0)
    return res


# 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(self, x1, x2):
    batch_size, n, num_dims = x1.size()
    _, m, _ = x2.size()

    if not num_dims == self.ard_num_dims:
        raise RuntimeError(
            "The SpectralMixtureKernel expected the input to have {} dimensionality "
            "(based on the ard_num_dims argument). Got {}.".format(
                self.ard_num_dims, num_dims
            )
        )
    if not batch_size == self.batch_size:
        raise RuntimeError(
            "The SpectralMixtureKernel expected the input to have a batch_size of {} "
            "(based on the batch_size argument). Got {}.".format(
                self.batch_size, batch_size
            )
        )

    # Expand x1 and x2 to account for the number of mixtures
    # Should make x1/x2 (b x k x n x d) for k mixtures
    x1_ = x1.unsqueeze(1)
    x2_ = x2.unsqueeze(1)

    # Compute distances - scaled by appropriate parameters
    x1_exp = x1_ * self.mixture_scales
    x2_exp = x2_ * self.mixture_scales
    x1_cos = x1_ * self.mixture_means
    x2_cos = x2_ * self.mixture_means

    # Create grids
    x1_exp_, x2_exp_ = self._create_input_grid(x1_exp, x2_exp)
    x1_cos_, x2_cos_ = self._create_input_grid(x1_cos, x2_cos)

    # Compute the exponential and cosine terms
    exp_term = (x1_exp_ - x2_exp_).pow_(2).mul_(-2 * math.pi**2)
    cos_term = (x1_cos_ - x2_cos_).mul_(2 * math.pi)
    res = exp_term.exp_() * cos_term.cos_()

    # Product omer dimensions
    res = res.prod(-1)

    # Sum over mixtures
    mixture_weights = self.mixture_weights
    while mixture_weights.dim() < res.dim():
        mixture_weights.unsqueeze_(-1)
    res = (res * mixture_weights).sum(1)
    return res