# Buggy code
def __init__(self, mean, covariance_matrix, validate_args=False, interleaved=True):
    """
    Constructs a multi-output multivariate Normal random variable, based on mean and covariance
    Can be multi-output multivariate, or a batch of multi-output multivariate Normal

    Passing a matrix mean corresponds to a multi-output multivariate Normal
    Passing a matrix mean corresponds to a batch of multivariate Normals

    Params:
        mean (:obj:`torch.tensor`): An `n x t` or batch `b x n x t` matrix of means for the MVN distribution.
        covar (:obj:`torch.tensor` or :obj:`gpytorch.lazy.LazyTensor`): An `nt x nt` or batch `b x nt x nt`
            covariance matrix of MVN distribution.
        validate_args (:obj:`bool`): If True, validate `mean` anad `covariance_matrix` arguments.
        interleaved (:obj:`bool`): If True, covariance matrix is interpreted as block-diagonal w.r.t.
            inter-task covariances for each observation. If False, it is interpreted as block-diagonal
            w.r.t. inter-observation covariance for each task.
    """
    if not torch.is_tensor(mean) and not isinstance(mean, LazyTensor):
        raise RuntimeError(
            "The mean of a MultitaskMultivariateNormal must be a Tensor or LazyTensor"
        )

    if not torch.is_tensor(covariance_matrix) and not isinstance(
        covariance_matrix, LazyTensor
    ):
        raise RuntimeError(
            "The covariance of a MultitaskMultivariateNormal must be a Tensor or LazyTensor"
        )

    if mean.dim() < 2:
        raise RuntimeError("mean should be a matrix or a batch matrix (batch mode)")

    self._output_shape = mean.shape
    # TODO: Instead of transpose / view operations, use a PermutationLazyTensor (see #539) to handle interleaving
    self._interleaved = interleaved
    if self._interleaved:
        mean_mvn = mean.reshape(*mean.shape[:-2], -1)
    else:
        mean_mvn = mean.transpose(-1, -2).reshape(*mean.shape[:-2], -1)
    super().__init__(
        mean=mean_mvn, covariance_matrix=covariance_matrix, validate_args=validate_args
    )


# 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, mean, covariance_matrix, validate_args=False, interleaved=True):
    if not torch.is_tensor(mean) and not isinstance(mean, LazyTensor):
        raise RuntimeError(
            "The mean of a MultitaskMultivariateNormal must be a Tensor or LazyTensor"
        )

    if not torch.is_tensor(covariance_matrix) and not isinstance(
        covariance_matrix, LazyTensor
    ):
        raise RuntimeError(
            "The covariance of a MultitaskMultivariateNormal must be a Tensor or LazyTensor"
        )

    if mean.dim() < 2:
        raise RuntimeError("mean should be a matrix or a batch matrix (batch mode)")

    self._output_shape = mean.shape
    # TODO: Instead of transpose / view operations, use a PermutationLazyTensor (see #539) to handle interleaving
    self._interleaved = interleaved
    if self._interleaved:
        mean_mvn = mean.reshape(*mean.shape[:-2], -1)
    else:
        mean_mvn = mean.transpose(-1, -2).reshape(*mean.shape[:-2], -1)
    super().__init__(
        mean=mean_mvn, covariance_matrix=covariance_matrix, validate_args=validate_args
    )