PyKeOps crashes when trying to use `torch.multiprocessing` with it

[GabrielMajeri]

I've been using PyKeOps successfully for a kernel ridge regression implementation, but when I tried to parallelize my code by using multiprocessing, I started to run into errors.

After some debugging, I've been able to find this minimal reproduction:

import torch
import torch.multiprocessing as mp

# Uncomment to trigger crash:
# from pykeops.torch import LazyTensor


def main() -> None:
    # Make sure we have at least one GPU with CUDA support
    assert torch.cuda.is_available()

    zeros = torch.zeros((64,), device="cuda")

    process = mp.Process(
        target=_worker_main,
        args=(zeros,),
    )
    process.start()

    process.join()


def _worker_main(tensor: torch.Tensor) -> None:
    print(tensor.ndim)
    del tensor

if __name__ == "__main__":
    main()

I'm using torch.multiprocessing to ensure that my tensors' CUDA memory buffers will automatically be shared across processes. However, as soon as I import PyKeOps, even before creating any LazyTensors, it leads to a crash:

.../site-packages/torch/multiprocessing/reductions.py", line 179, in rebuild_cuda_tensor
    torch.cuda._lazy_init()
    ~~~~~~~~~~~~~~~~~~~~~^^
.../site-packages/torch/cuda/__init__.py", line 478, in _lazy_init
    torch._C._cuda_init()
    ~~~~~~~~~~~~~~~~~~~^^
RuntimeError: CUDA driver initialization failed, you might not have a CUDA gpu.

I've been unable to pinpoint the exact nature of the error, but if I were to guess, it's probably due to PyTorch trying to initialize CUDA, but KeOps already initialized it?