# Buggy code
def __init__(self, num_inducing_points, batch_shape=torch.Size([]), **kwargs):
    """
    Args:
        num_inducing_points (int): Size of the variational distribution. This implies that the variational mean
            should be this size, and the variational covariance matrix should have this many rows and columns.
        batch_shape (torch.Size, optional): Specifies an optional batch
            size for the variational parameters. This is useful for example
            when doing additive variational inference.
    """
    batch_shape = _deprecate_kwarg_with_transform(
        kwargs, "batch_size", "batch_shape", batch_shape, lambda n: torch.Size([n])
    )
    super(VariationalDistribution, self).__init__()
    mean_init = torch.zeros(num_inducing_points)
    covar_init = torch.eye(num_inducing_points, num_inducing_points)
    mean_init = mean_init.repeat(*batch_shape, 1)
    covar_init = covar_init.repeat(*batch_shape, 1, 1)

    self.register_parameter(
        name="variational_mean", parameter=torch.nn.Parameter(mean_init)
    )
    self.register_parameter(
        name="chol_variational_covar", parameter=torch.nn.Parameter(covar_init)
    )


# 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 __init__(
    self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=1e-3, **kwargs
):
    super().__init__(
        num_inducing_points=num_inducing_points,
        batch_shape=batch_shape,
        mean_init_std=mean_init_std,
    )
    mean_init = torch.zeros(num_inducing_points)
    covar_init = torch.eye(num_inducing_points, num_inducing_points)
    mean_init = mean_init.repeat(*batch_shape, 1)
    covar_init = covar_init.repeat(*batch_shape, 1, 1)

    self.register_parameter(
        name="variational_mean", parameter=torch.nn.Parameter(mean_init)
    )
    self.register_parameter(
        name="chol_variational_covar", parameter=torch.nn.Parameter(covar_init)
    )