# Buggy code
def _getitem(self, row_index, col_index, *batch_indices):
    """
    Supports subindexing of the matrix this LazyTensor represents.

    The indices passed into this method will either be:
        Tensor indices
        Slices

    ..note::
        LazyTensor.__getitem__ uses this as a helper method. If you are writing your own custom LazyTensor,
        override this method rather than __getitem__ (so that you don't have to repeat the extra work)

    ..note::
        This method is used internally by the related function :func:`~gpytorch.lazy.LazyTensor.__getitem__`,
        which does some additional work. Calling this method directly is discouraged.

    This method has a number of restrictions on the type of arguments that are passed in to reduce
    the complexity of __getitem__ calls in PyTorch. In particular:
        - This method only accepts slices and tensors for the row/column indices (no ints)
        - The row and column dimensions don't dissapear (e.g. from Tensor indexing). These cases are
          handled by the `_getindices` method

    Args:
        :attr:`row_index` (slice, Tensor):
            Index for the row of the LazyTensor
        :attr:`col_index` (slice, Tensor):
            Index for the col of the LazyTensor
        :attr:`batch_indices` (tuple of slice, int, Tensor):
            Indices for the batch dimensions

    Returns:
        `LazyTensor`
    """
    # Special case: if both row and col are not indexed, then we are done
    if row_index is _noop_index and col_index is _noop_index:
        if len(batch_indices):
            components = [component[batch_indices] for component in self._args]
            res = self.__class__(*components, **self._kwargs)
            return res
        else:
            return self

    # Normal case: we have to do some processing on either the rows or columns
    # We will handle this through "interpolation"
    row_interp_indices = torch.arange(
        0, self.size(-2), dtype=torch.long, device=self.device
    ).view(-1, 1)
    row_interp_indices = row_interp_indices.expand(*self.batch_shape, -1, 1)
    row_interp_values = torch.tensor(
        1.0, dtype=self.dtype, device=self.device
    ).expand_as(row_interp_indices)

    col_interp_indices = torch.arange(
        0, self.size(-1), dtype=torch.long, device=self.device
    ).view(-1, 1)
    col_interp_indices = col_interp_indices.expand(*self.batch_shape, -1, 1)
    col_interp_values = torch.tensor(
        1.0, dtype=self.dtype, device=self.device
    ).expand_as(col_interp_indices)

    # Construct interpolated LazyTensor
    from . import InterpolatedLazyTensor

    res = InterpolatedLazyTensor(
        self,
        row_interp_indices,
        row_interp_values,
        col_interp_indices,
        col_interp_values,
    )
    return res._getitem(row_index, col_index, *batch_indices)


# Traceback (IndexError)
Traceback (most recent call last):
File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py", line 88, in _getitem
x1 = x1[(*batch_indices, row_index, dim_index)]
IndexError: too many indices for tensor of dimension 2

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/Users/johnsonzhong/Research/meta_contact/simulation/temp.py", line 70, in <module>
lower, upper = predictions.confidence_region()
File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/distributions/multivariate_normal.py", line 80, in confidence_region
std2 = self.stddev.mul_(2)
File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/torch/distributions/distribution.py", line 111, in stddev
return self.variance.sqrt()
File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/distributions/multitask_multivariate_normal.py", line 219, in variance
var = super().variance
File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/distributions/multivariate_normal.py", line 189, in variance
diag = self.lazy_covariance_matrix.diag()
File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/sum_lazy_tensor.py", line 96, in diag
return sum(lazy_tensor.diag().contiguous() for lazy_tensor in self.lazy_tensors)
File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/sum_lazy_tensor.py", line 96, in <genexpr>
return sum(lazy_tensor.diag().contiguous() for lazy_tensor in self.lazy_tensors)
File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/interpolated_lazy_tensor.py", line 387, in diag
return super(InterpolatedLazyTensor, self).diag()
File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py", line 839, in diag
return self[..., row_col_iter, row_col_iter]
File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py", line 1716, in __getitem__
res = self._get_indices(row_index, col_index, *batch_indices)
File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/interpolated_lazy_tensor.py", line 110, in _get_indices
*[batch_index.view(*batch_index.shape, 1, 1) for batch_index in batch_indices],
File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/block_interleaved_lazy_tensor.py", line 58, in _get_indices
res = self.base_lazy_tensor._get_indices(row_index, col_index, *batch_indices, row_index_block)
File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py", line 294, in _get_indices
base_lazy_tensor = self._getitem(_noop_index, _noop_index, *batch_indices)._expand_batch(final_shape)
File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py", line 97, in _getitem
f"Attempting to tensor index a non-batch matrix's batch dimensions. "
RuntimeError: Attempting to tensor index a non-batch matrix's batch dimensions. Got batch index {batch_indices} but my shape was {self.shape}

# Fixed code
def _getitem(self, row_index, col_index, *batch_indices):
    """
    Supports subindexing of the matrix this LazyTensor represents.

    The indices passed into this method will either be:
        Tensor indices
        Slices

    ..note::
        LazyTensor.__getitem__ uses this as a helper method. If you are writing your own custom LazyTensor,
        override this method rather than __getitem__ (so that you don't have to repeat the extra work)

    ..note::
        This method is used internally by the related function :func:`~gpytorch.lazy.LazyTensor.__getitem__`,
        which does some additional work. Calling this method directly is discouraged.

    This method has a number of restrictions on the type of arguments that are passed in to reduce
    the complexity of __getitem__ calls in PyTorch. In particular:
        - This method only accepts slices and tensors for the row/column indices (no ints)
        - The row and column dimensions don't dissapear (e.g. from Tensor indexing). These cases are
          handled by the `_getindices` method

    Args:
        :attr:`row_index` (slice, Tensor):
            Index for the row of the LazyTensor
        :attr:`col_index` (slice, Tensor):
            Index for the col of the LazyTensor
        :attr:`batch_indices` (tuple of slice, int, Tensor):
            Indices for the batch dimensions

    Returns:
        `LazyTensor`
    """
    # Special case: if both row and col are not indexed, then we are done
    if _is_noop_index(row_index) and _is_noop_index(col_index):
        if len(batch_indices):
            components = [component[batch_indices] for component in self._args]
            res = self.__class__(*components, **self._kwargs)
            return res
        else:
            return self

    # Normal case: we have to do some processing on either the rows or columns
    # We will handle this through "interpolation"
    row_interp_indices = torch.arange(
        0, self.size(-2), dtype=torch.long, device=self.device
    ).view(-1, 1)
    row_interp_indices = row_interp_indices.expand(*self.batch_shape, -1, 1)
    row_interp_values = torch.tensor(
        1.0, dtype=self.dtype, device=self.device
    ).expand_as(row_interp_indices)

    col_interp_indices = torch.arange(
        0, self.size(-1), dtype=torch.long, device=self.device
    ).view(-1, 1)
    col_interp_indices = col_interp_indices.expand(*self.batch_shape, -1, 1)
    col_interp_values = torch.tensor(
        1.0, dtype=self.dtype, device=self.device
    ).expand_as(col_interp_indices)

    # Construct interpolated LazyTensor
    from . import InterpolatedLazyTensor

    res = InterpolatedLazyTensor(
        self,
        row_interp_indices,
        row_interp_values,
        col_interp_indices,
        col_interp_values,
    )
    return res._getitem(row_index, col_index, *batch_indices)