IndexError when using torch_cuda backend

[const7]

Thanks for your wonderful tool! I recently used himalaya for voxel-wise encoding but got the IndexError when using torch_cuda backend.

Minimal reproducible code

The example code from your document: Multiple-kernel ridge with scikit-learn API

Log

[                                        ] 0% | 0.00 sec | 100 random sampling with cv | Traceback (most recent call last):
  File "C:\Users\kun\Downloads\plot_mkr_1_sklearn_api.py", line 108, in <module>
    pipe.fit(X_train, Y_train)
  File "C:\Users\kun\miniconda3\envs\naremo\lib\site-packages\sklearn\base.py", line 1473, in wrapper
    return fit_method(estimator, *args, **kwargs)
  File "C:\Users\kun\miniconda3\envs\naremo\lib\site-packages\sklearn\pipeline.py", line 473, in fit
    self._final_estimator.fit(Xt, y, **last_step_params["fit"])
  File "C:\Users\kun\miniconda3\envs\naremo\lib\site-packages\himalaya\backend\_utils.py", line 97, in wrapper        
    return func(*args, **kwargs)
  File "C:\Users\kun\miniconda3\envs\naremo\lib\site-packages\himalaya\kernel_ridge\_sklearn_api.py", line 935, in fit
    tmp = self._call_solver(Ks=Ks, Y=y, cv=cv, return_weights="dual",
  File "C:\Users\kun\miniconda3\envs\naremo\lib\site-packages\himalaya\kernel_ridge\_sklearn_api.py", line 59, in _call_solver
    return function(**direct_params, **solver_params)
  File "C:\Users\kun\miniconda3\envs\naremo\lib\site-packages\himalaya\kernel_ridge\_random_search.py", line 222, in solve_multiple_kernel_ridge_random_search
    Ktrain, Ktest = K[train[:, None], train], K[test[:, None], train]
IndexError: tensors used as indices must be long, byte or bool tensors

Environment

Here is my torch information:

PyTorch version: 1.13.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 Pro
GCC version: (GCC) 11.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.10.14 | packaged by Anaconda, Inc. | (main, May  6 2024, 19:44:50) [MSC v.1916 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.22631-SP0
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4060 Ti
Nvidia driver version: 560.70
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] pytorchvideo==0.1.5
[pip3] torch==1.13.1+cu117
[pip3] torchaudio==0.13.1+cu117
[pip3] torchvision==0.14.1+cu117
[conda] blas                      1.0                         mkl
[conda] mkl                       2021.4.0                 pypi_0    pypi
[conda] mkl-service               2.4.0           py310h2bbff1b_1
[conda] mkl_fft                   1.3.8           py310h2bbff1b_0
[conda] mkl_random                1.2.4           py310h59b6b97_0
[conda] numpy                     1.26.4          py310h055cbcc_0
[conda] numpy-base                1.26.4          py310h65a83cf_0
[conda] pytorchvideo              0.1.5                    pypi_0    pypi
[conda] torch                     1.13.1+cu117             pypi_0    pypi
[conda] torchaudio                0.13.1+cu117             pypi_0    pypi
[conda] torchvision               0.14.1+cu117             pypi_0    pypi

others:

Package            Version
himalaya           0.4.6
scikit-learn        1.5.1
numpy               1.26.4

Please let me know if you need more information. Thank you!