# Buggy code
def marginal(self, function_dist, *params, **kwargs):
    name_prefix = kwargs.get("name_prefix", "")
    num_samples = settings.num_likelihood_samples.value()
    with pyro.plate(
        name_prefix + ".num_particles_vectorized",
        num_samples,
        dim=(-self.max_plate_nesting - 1),
    ):
        function_samples_shape = torch.Size(
            [num_samples]
            + [1] * (self.max_plate_nesting - len(function_dist.batch_shape) - 1)
        )
        function_samples = function_dist(function_samples_shape)
        if self.training:
            return self(function_samples, *params, **kwargs)
        else:
            guide_trace = pyro.poutine.trace(self.guide).get_trace(*params, **kwargs)
            marginal_fn = functools.partial(self.__call__, function_samples)
            return pyro.poutine.replay(marginal_fn, trace=guide_trace)(
                *params, **kwargs
            )


# Traceback (AttributeError)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-6-c955f39ee560> in <module>
4 with torch.no_grad():
5     for x_batch, y_batch in test_loader:
----> 6         preds = model(x_batch)
7         means = torch.cat([means, preds.mean.cpu()])
8 means = means[1:]

~/miniconda3/envs/mbrl/lib/python3.7/site-packages/gpytorch/models/abstract_variational_gp.py in __call__(self, inputs, **kwargs)
20             inputs = inputs.unsqueeze(-1)
21
---> 22         return self.variational_strategy(inputs)

~/miniconda3/envs/mbrl/lib/python3.7/site-packages/gpytorch/variational/variational_strategy.py in __call__(self, x)
228                 self._memoize_cache = dict()
229
--> 230         return super(VariationalStrategy, self).__call__(x)

~/miniconda3/envs/mbrl/lib/python3.7/site-packages/gpytorch/module.py in __call__(self, *inputs, **kwargs)
20
21     def __call__(self, *inputs, **kwargs):
---> 22         outputs = self.forward(*inputs, **kwargs)
23         if isinstance(outputs, list):
24             return [_validate_module_outputs(output) for output in outputs]

~/miniconda3/envs/mbrl/lib/python3.7/site-packages/gpytorch/variational/variational_strategy.py in forward(self, x)
214                     induc_induc_covar.inv_matmul(induc_data_covar)
215                 )
--> 216                 data_covariance = data_data_covar + neg_induc_data_data_covar
217             predictive_covar = PsdSumLazyTensor(predictive_covar, data_covariance)
218

~/miniconda3/envs/mbrl/lib/python3.7/site-packages/gpytorch/lazy/sum_lazy_tensor.py in __add__(self, other)
75             return SumLazyTensor(*(list(self.lazy_tensors) + [other]))
76         else:
---> 77             raise AttributeError("other must be a LazyTensor")
78
79     def diag(self):

AttributeError: other must be a LazyTensor

# Fixed code
def marginal(self, function_dist, *args, **kwargs):
    r"""
    Computes a predictive distribution :math:`p(y^* | \mathbf x^*)` given either a posterior
    distribution :math:`p(\mathbf f | \mathcal D, \mathbf x)` or a
    prior distribution :math:`p(\mathbf f|\mathbf x)` as input.

    With both exact inference and variational inference, the form of
    :math:`p(\mathbf f|\mathcal D, \mathbf x)` or :math:`p(\mathbf f|
    \mathbf x)` should usually be Gaussian. As a result, :attr:`function_dist`
    should usually be a :obj:`~gpytorch.distributions.MultivariateNormal` specified by the mean and
    (co)variance of :math:`p(\mathbf f|...)`.

    Args:
        :attr:`function_dist` (:class:`~gpytorch.distributions.MultivariateNormal`)
            Distribution for :math:`f(x)`.
        :attr:`args`, :attr:`kwargs`
            Passed to the `forward` function

    Returns:
        Distribution object (the marginal distribution, or samples from it)
    """
    return super().marginal(function_dist, *args, **kwargs)