# Buggy code
def __call__(self, inputs, are_samples=False, **kwargs):
    """
    Forward data through this hidden GP layer. The output is a MultitaskMultivariateNormal distribution
    (or MultivariateNormal distribution is output_dims=None).

    If the input is >=2 dimensional Tensor (e.g. `n x d`), we pass the input through each hidden GP,
    resulting in a `n x h` multitask Gaussian distribution (where all of the `h` tasks represent an
    output dimension and are independent from one another).  We then draw `s` samples from these Gaussians,
    resulting in a `s x n x h` MultitaskMultivariateNormal distribution.

    If the input is a >=3 dimensional Tensor, and the `are_samples=True` kwarg is set, then we assume that
    the outermost batch dimension is a samples dimension. The output will have the same number of samples.
    For example, a `s x b x n x d` input will result in a `s x b x n x h` MultitaskMultivariateNormal distribution.

    The goal of these last two points is that if you have a tensor `x` that is `n x d`, then:
        >>> hidden_gp2(hidden_gp(x))

    will just work, and return a tensor of size `s x n x h2`, where `h2` is the output dimensionality of
    hidden_gp2. In this way, hidden GP layers are easily composable.
    """
    deterministic_inputs = not are_samples
    if isinstance(inputs, MultitaskMultivariateNormal):
        inputs = torch.distributions.Normal(
            loc=inputs.mean, scale=inputs.variance.sqrt()
        ).rsample()
        deterministic_inputs = False

    if settings.debug.on():
        if not torch.is_tensor(inputs):
            raise ValueError(
                "`inputs` should either be a MultitaskMultivariateNormal or a Tensor, got "
                f"{inputs.__class__.__Name__}"
            )

        if inputs.size(-1) != self.input_dims:
            raise RuntimeError(
                f"Input shape did not match self.input_dims. Got total feature dims [{inputs.size(-1)}],"
                f" expected [{self.input_dims}]"
            )

    # Repeat the input for all possible outputs
    if self.output_dims is not None:
        inputs = inputs.unsqueeze(-3)
        inputs = inputs.expand(*inputs.shape[:-3], self.output_dims, *inputs.shape[-2:])

    # Now run samples through the GP
    output = AbstractVariationalGP.__call__(self, inputs)
    if self.output_dims is not None:
        mean = output.loc.transpose(-1, -2)
        covar = BlockDiagLazyTensor(output.lazy_covariance_matrix, block_dim=-3)
        output = MultitaskMultivariateNormal(mean, covar, interleaved=False)

    # Maybe expand inputs?
    if deterministic_inputs:
        output = output.expand(
            torch.Size([settings.num_likelihood_samples.value()]) + output.batch_shape
        )

    return output


# 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, inputs, are_samples=False, **kwargs):
    """
    Forward data through this hidden GP layer. The output is a MultitaskMultivariateNormal distribution
    (or MultivariateNormal distribution is output_dims=None).

    If the input is >=2 dimensional Tensor (e.g. `n x d`), we pass the input through each hidden GP,
    resulting in a `n x h` multitask Gaussian distribution (where all of the `h` tasks represent an
    output dimension and are independent from one another).  We then draw `s` samples from these Gaussians,
    resulting in a `s x n x h` MultitaskMultivariateNormal distribution.

    If the input is a >=3 dimensional Tensor, and the `are_samples=True` kwarg is set, then we assume that
    the outermost batch dimension is a samples dimension. The output will have the same number of samples.
    For example, a `s x b x n x d` input will result in a `s x b x n x h` MultitaskMultivariateNormal distribution.

    The goal of these last two points is that if you have a tensor `x` that is `n x d`, then:
        >>> hidden_gp2(hidden_gp(x))

    will just work, and return a tensor of size `s x n x h2`, where `h2` is the output dimensionality of
    hidden_gp2. In this way, hidden GP layers are easily composable.
    """
    deterministic_inputs = not are_samples
    if isinstance(inputs, MultitaskMultivariateNormal):
        inputs = torch.distributions.Normal(
            loc=inputs.mean, scale=inputs.variance.sqrt()
        ).rsample()
        deterministic_inputs = False

    if settings.debug.on():
        if not torch.is_tensor(inputs):
            raise ValueError(
                "`inputs` should either be a MultitaskMultivariateNormal or a Tensor, got "
                f"{inputs.__class__.__Name__}"
            )

        if inputs.size(-1) != self.input_dims:
            raise RuntimeError(
                f"Input shape did not match self.input_dims. Got total feature dims [{inputs.size(-1)}],"
                f" expected [{self.input_dims}]"
            )

    # Repeat the input for all possible outputs
    if self.output_dims is not None:
        inputs = inputs.unsqueeze(-3)
        inputs = inputs.expand(*inputs.shape[:-3], self.output_dims, *inputs.shape[-2:])

    # Now run samples through the GP
    output = ApproximateGP.__call__(self, inputs)
    if self.output_dims is not None:
        mean = output.loc.transpose(-1, -2)
        covar = BlockDiagLazyTensor(output.lazy_covariance_matrix, block_dim=-3)
        output = MultitaskMultivariateNormal(mean, covar, interleaved=False)

    # Maybe expand inputs?
    if deterministic_inputs:
        output = output.expand(
            torch.Size([settings.num_likelihood_samples.value()]) + output.batch_shape
        )

    return output