# Buggy code
def __call__(self, x):
    if not self.variational_params_initialized.item():
        self.initialize_variational_dist()
        self.variational_params_initialized.fill_(1)
    if self.training:
        if hasattr(self, "_memoize_cache"):
            delattr(self, "_memoize_cache")
            self._memoize_cache = dict()

    return super(VariationalStrategy, self).__call__(x)


# 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 __call__(self, x, prior=False):
    if not self.updated_strategy.item() and not prior:
        with torch.no_grad():
            # Get unwhitened p(u)
            prior_function_dist = self(self.inducing_points, prior=True)
            prior_mean = prior_function_dist.loc
            L = self._cholesky_factor(
                prior_function_dist.lazy_covariance_matrix.add_jitter()
            )

            # Temporarily turn off noise that's added to the mean
            orig_mean_init_std = self._variational_distribution.mean_init_std
            self._variational_distribution.mean_init_std = 0.0

            # Change the variational parameters to be whitened
            variational_dist = self.variational_distribution
            whitened_mean = (
                torch.triangular_solve(
                    (variational_dist.loc - prior_mean).unsqueeze(-1).double(),
                    L,
                    upper=False,
                )[0]
                .squeeze(-1)
                .to(variational_dist.loc.dtype)
            )
            whitened_covar = RootLazyTensor(
                torch.triangular_solve(
                    variational_dist.lazy_covariance_matrix.root_decomposition()
                    .root.evaluate()
                    .double(),
                    L,
                    upper=False,
                )[0].to(variational_dist.loc.dtype)
            )
            whitened_variational_distribution = variational_dist.__class__(
                whitened_mean, whitened_covar
            )
            self._variational_distribution.initialize_variational_distribution(
                whitened_variational_distribution
            )

            # Reset the random noise parameter of the model
            self._variational_distribution.mean_init_std = orig_mean_init_std

            # Reset the cache
            if hasattr(self, "_memoize_cache"):
                delattr(self, "_memoize_cache")
                self._memoize_cache = dict()

            # Mark that we have updated the variational strategy
            self.updated_strategy.fill_(True)

    return super().__call__(x, prior=prior)