[Bug] Spurious error raised while indexing a RBFKernelGrad LazyTensor #1554
Description
🐛 Bug
While indexing RBFKernelGrad, I get a spurious error
To reproduce
import torch
import gpytorch
x=torch.tensor([0.,1.,2.])
index=torch.tensor([True,False,True])
full_kernel = gpytorch.kernels.RBFKernel()(x)
full_kernel[:,index]
full_kernel[:,index][index,:]
print('Slicing worked fine for RBF kernel')
full_kernel = gpytorch.kernels.RBFKernelGrad()(x)
index=torch.cat([index,index])
full_kernel[:,index]
print('However, the next line raises an error for RBFGrad kernel')
full_kernel[:,index][index,:]** Stack trace/error message **
Slicing worked fine for RBF kernel
However, it raises an error for RBFGrad kernel
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
12 full_kernel[:,index]
13 print('However, it raises an error for RBFGrad kernel')
---> 14 full_kernel[:,index][index,:]
/opt/anaconda3/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py in __getitem__(self, index)
1730 raise RuntimeError(
1731 "{}.__getitem__ failed! Expected a final shape of size {}, got {}. This is a bug with GPyTorch, "
-> 1732 "or your custom LazyTensor.".format(self.__class__.__name__, expected_shape, res.shape)
1733 )
1734
RuntimeError: NonLazyTensor.__getitem__ failed! Expected a final shape of size torch.Size([6, 4]), got torch.Size([4, 4]). This is a bug with GPyTorch, or your custom LazyTensor.
Expected Behavior
A [4,4] tensor size is expected, so I do not see how the Expected a final shape of size torch.Size([6, 4]) is computed. Therefore, I believe the error should not be raised.
System information
- GPyTorch Version 1.1.1
- PyTorch Version 1.8.0
- Computer OS MacOS 11.2.3