[Bug] Indexing LazyEvaluatedKernelTensor on batch dimension with index greater than 0 gives IndexError

[tristancinquin]

🐛 Bug

Hi, I am trying to index the batch dimension of a LazyEvaluatedKernelTensor but this leads to an IndexError whenever the index is greater than 0.

To reproduce

Code snippet to reproduce

import torch, gpytorch

x = torch.linspace(-1, 1, 100).reshape(-1, 1)
# kernel with a batch output size of 2
kernel = gpytorch.kernels.RBFKernel(ard_num_dims=1, active_dims=[0], batch_shape=torch.Size([2]))
print(kernel(x)) # output has torch.Size([2, 100, 100])
print(kernel(x)[0,:,:]) # output has torch.Size([100, 100])
print(kernel(x)[1,:,:]) # IndexError

Stack trace/error message

Traceback (most recent call last):
  File "/Users/tristancinquin/anaconda3/envs/bonntorch/lib/python3.12/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py", line 228, in _getitem
    new_kernel = self.kernel.__getitem__(batch_indices)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tristancinquin/anaconda3/envs/bonntorch/lib/python3.12/site-packages/gpytorch/kernels/kernel.py", line 579, in __getitem__
    new_buffr.data = new_buffr.__getitem__(index)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
IndexError: index 1 is out of bounds for dimension 0 with size 1

Expected Behavior

print(kernel(x)[1,:,:]) should not raise an IndexError and should return a LazyEvaluatedKernelTensor of size (100, 100).

System information

Please complete the following information:

• GPyTorch: 1.12
• PyTorch: 2.4.1
• Python: 3.12.0
• Computer OS: macOS Sonoma 14.6.1