# Buggy code
def root_inv_decomposition(self, initial_vectors=None, test_vectors=None):
    """
    Returns a (usually low-rank) root decomposotion lazy tensor of a PSD matrix.
    This can be used for sampling from a Gaussian distribution, or for obtaining a
    low-rank version of a matrix
    """
    from .root_lazy_tensor import RootLazyTensor

    if not self.is_square:
        raise RuntimeError(
            "root_inv_decomposition only operates on (batches of) square (symmetric) LazyTensors. "
            "Got a {} of size {}.".format(self.__class__.__name__, self.size())
        )

    if initial_vectors is not None:
        if self.dim() == 2 and initial_vectors.dim() == 1:
            if self.shape[-1] != initial_vectors.numel():
                raise RuntimeError(
                    "LazyTensor (size={}) cannot be multiplied with initial_vectors (size={}).".format(
                        self.shape, initial_vectors.shape
                    )
                )
        elif self.dim() != initial_vectors.dim():
            raise RuntimeError(
                "LazyTensor (size={}) and initial_vectors (size={}) should have the same number "
                "of dimensions.".format(self.shape, initial_vectors.shape)
            )
        elif (
            self.batch_shape != initial_vectors.shape[:-2]
            or self.shape[-1] != initial_vectors.shape[-2]
        ):
            raise RuntimeError(
                "LazyTensor (size={}) cannot be multiplied with initial_vectors (size={}).".format(
                    self.shape, initial_vectors.shape
                )
            )

    roots, inv_roots = RootDecomposition(
        self.representation_tree(),
        max_iter=self.root_decomposition_size(),
        dtype=self.dtype,
        device=self.device,
        batch_shape=self.batch_shape,
        matrix_shape=self.matrix_shape,
        root=True,
        inverse=True,
        initial_vectors=initial_vectors,
    )(*self.representation())

    if initial_vectors is not None and initial_vectors.size(-1) > 1:
        self._memoize_cache["root_decomposition"] = RootLazyTensor(roots[0])
    else:
        self._memoize_cache["root_decomposition"] = RootLazyTensor(roots)

    # Choose the best of the inv_roots, if there were more than one initial vectors
    if initial_vectors is not None and initial_vectors.size(-1) > 1:
        num_probes = initial_vectors.size(-1)
        test_vectors = test_vectors.unsqueeze(0)

        # Compute solves
        solves = inv_roots.matmul(inv_roots.transpose(-1, -2).matmul(test_vectors))

        # Compute self * solves
        solves = (
            solves.permute(*range(1, self.dim() + 1), 0)
            .contiguous()
            .view(*self.batch_shape, self.matrix_shape[-1], -1)
        )
        mat_times_solves = self.matmul(solves)
        mat_times_solves = mat_times_solves.view(
            *self.batch_shape, self.matrix_shape[-1], -1, num_probes
        ).permute(-1, *range(0, self.dim()))

        # Compute residuals
        residuals = (mat_times_solves - test_vectors).norm(2, dim=-2)
        residuals = residuals.view(residuals.size(0), -1).sum(-1)

        # Choose solve that best fits
        _, best_solve_index = residuals.min(0)
        inv_root = inv_roots[best_solve_index].squeeze(0)

    else:
        inv_root = inv_roots

    return RootLazyTensor(inv_root)


# Traceback (IndexError)
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-1-05bfb3c2646a> in <module>
27 # this throws the error
28 with gpytorch.settings.fast_pred_var():
---> 29     model(torch.rand(100, 2))

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/models/exact_gp.py in __call__(self, *args, **kwargs)
262
263             predictive_mean = self.prediction_strategy.exact_predictive_mean(test_mean, test_train_covar)
--> 264             predictive_covar = self.prediction_strategy.exact_predictive_covar(test_test_covar, test_train_covar)
265
266             if num_tasks > 1:

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/models/exact_prediction_strategies.py in exact_predictive_covar(self, test_test_covar, test_train_covar)
324             return test_test_covar + MatmulLazyTensor(test_train_covar, covar_correction_rhs)
325
--> 326         precomputed_cache = self.covar_cache
327         covar_inv_quad_form_root = self._exact_predictive_covar_inv_quad_form_root(precomputed_cache,
328                                                                                    test_train_covar)

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/utils/memoize.py in g(self, *args, **kwargs)
15         cache_name = name if name is not None else method
16         if cache_name not in self._memoize_cache:
---> 17             self._memoize_cache[cache_name] = method(self, *args, **kwargs)
18         return self._memoize_cache[cache_name]
19

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/models/exact_prediction_strategies.py in covar_cache(self)
265             train_train_covar_inv_root = train_train_covar[0].root_inv_decomposition().root.evaluate()
266         else:
--> 267             train_train_covar_inv_root = train_train_covar.root_inv_decomposition().root.evaluate()
268
269         return self._exact_predictive_covar_inv_quad_form_cache(train_train_covar_inv_root, self._last_test_train_covar)

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/utils/memoize.py in g(self, *args, **kwargs)
15         cache_name = name if name is not None else method
16         if cache_name not in self._memoize_cache:
---> 17             self._memoize_cache[cache_name] = method(self, *args, **kwargs)
18         return self._memoize_cache[cache_name]
19

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py in root_inv_decomposition(self, initial_vectors, test_vectors)
1088             inverse=True,
1089             initial_vectors=initial_vectors,
-> 1090         )(*self.representation())
1091
1092         if initial_vectors is not None and initial_vectors.size(-1) > 1:

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/functions/_root_decomposition.py in forward(self, *matrix_args)
51             matrix_shape=self.matrix_shape,
52             batch_shape=self.batch_shape,
---> 53             init_vecs=self.initial_vectors,
54         )
55

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/utils/lanczos.py in lanczos_tridiag(matmul_closure, max_iter, dtype, device, matrix_shape, batch_shape, init_vecs, num_init_vecs, tol)
80     # Copy over alpha_0 and beta_0 to t_mat
81     t_mat[0, 0].copy_(alpha_0)
---> 82     t_mat[0, 1].copy_(beta_0)
83     t_mat[1, 0].copy_(beta_0)
84

IndexError: index 1 is out of bounds for dimension 0 with size 1

# Fixed code
def root_inv_decomposition(self, initial_vectors=None, test_vectors=None):
    """
    Returns a (usually low-rank) root decomposotion lazy tensor of a PSD matrix.
    This can be used for sampling from a Gaussian distribution, or for obtaining a
    low-rank version of a matrix
    """
    from .root_lazy_tensor import RootLazyTensor

    if self.shape[-2:].numel() == 1:
        return RootLazyTensor(1 / self.evaluate().sqrt())

    if not self.is_square:
        raise RuntimeError(
            "root_inv_decomposition only operates on (batches of) square (symmetric) LazyTensors. "
            "Got a {} of size {}.".format(self.__class__.__name__, self.size())
        )

    if initial_vectors is not None:
        if self.dim() == 2 and initial_vectors.dim() == 1:
            if self.shape[-1] != initial_vectors.numel():
                raise RuntimeError(
                    "LazyTensor (size={}) cannot be multiplied with initial_vectors (size={}).".format(
                        self.shape, initial_vectors.shape
                    )
                )
        elif self.dim() != initial_vectors.dim():
            raise RuntimeError(
                "LazyTensor (size={}) and initial_vectors (size={}) should have the same number "
                "of dimensions.".format(self.shape, initial_vectors.shape)
            )
        elif (
            self.batch_shape != initial_vectors.shape[:-2]
            or self.shape[-1] != initial_vectors.shape[-2]
        ):
            raise RuntimeError(
                "LazyTensor (size={}) cannot be multiplied with initial_vectors (size={}).".format(
                    self.shape, initial_vectors.shape
                )
            )

    roots, inv_roots = RootDecomposition(
        self.representation_tree(),
        max_iter=self.root_decomposition_size(),
        dtype=self.dtype,
        device=self.device,
        batch_shape=self.batch_shape,
        matrix_shape=self.matrix_shape,
        root=True,
        inverse=True,
        initial_vectors=initial_vectors,
    )(*self.representation())

    if initial_vectors is not None and initial_vectors.size(-1) > 1:
        self._memoize_cache["root_decomposition"] = RootLazyTensor(roots[0])
    else:
        self._memoize_cache["root_decomposition"] = RootLazyTensor(roots)

    # Choose the best of the inv_roots, if there were more than one initial vectors
    if initial_vectors is not None and initial_vectors.size(-1) > 1:
        num_probes = initial_vectors.size(-1)
        test_vectors = test_vectors.unsqueeze(0)

        # Compute solves
        solves = inv_roots.matmul(inv_roots.transpose(-1, -2).matmul(test_vectors))

        # Compute self * solves
        solves = (
            solves.permute(*range(1, self.dim() + 1), 0)
            .contiguous()
            .view(*self.batch_shape, self.matrix_shape[-1], -1)
        )
        mat_times_solves = self.matmul(solves)
        mat_times_solves = mat_times_solves.view(
            *self.batch_shape, self.matrix_shape[-1], -1, num_probes
        ).permute(-1, *range(0, self.dim()))

        # Compute residuals
        residuals = (mat_times_solves - test_vectors).norm(2, dim=-2)
        residuals = residuals.view(residuals.size(0), -1).sum(-1)

        # Choose solve that best fits
        _, best_solve_index = residuals.min(0)
        inv_root = inv_roots[best_solve_index].squeeze(0)

    else:
        inv_root = inv_roots

    return RootLazyTensor(inv_root)