Numba cuda and pytorch distribution may conflict

I’m trying to use numba and pytorch distribution simultaneously. When I create new tensor on GPU, I got cuda initialization error:
Here is my code:

import torch
import numpy as np
import torch.distributed as dist
from torch.multiprocessing import Process
from numba import cuda
import os
def func(idx):
    a = torch.randn([10, 10, 10], device=‘cuda’)
return
def init_process(idx, size, fn, backend=‘NCCL’):
    os.environ[‘MASTER_ADDR’] = ‘127.0.0.1’
    os.environ[‘MASTER_PORT’] = ‘29500’
    dist.init_process_group(backend, rank=idx, world_size=size)
    fn(idx)
    return
if  **name**  == “ **main** ”:
    processes = []
    for idx in range(2):
        p = Process(target=init_process, args=(idx, 2, func))
        p.start()
        processes.append(p)
    for p in processes:
        p.join()

Here is error report:

THCudaCheck FAIL file=/pytorch/aten/src/THC/THCGeneral.cpp line=47 error=3 : initialization error
Process Process-2:
Traceback (most recent call last):
File “/home/sss/anaconda3/envs/torch_new/lib/python3.8/multiprocessing/process.py”, line 315, in _bootstrap
self.run()
File “/home/sss/anaconda3/envs/torch_new/lib/python3.8/multiprocessing/process.py”, line 108, in run
self._target(*self._args, **self._kwargs)
File “/home/sss/Desktop/Experiment/test.py”, line 14, in init_process
fn(idx)
File “/home/sss/Desktop/Experiment/test.py”, line 8, in func
a = torch.randn([10, 10, 10], device=‘cuda’)
File “/home/sss/anaconda3/envs/torch_new/lib/python3.8/site-packages/torch/cuda/ **init** .py”, line 190, in _lazy_init
torch._C._cuda_init()
RuntimeError: cuda runtime error (3) : initialization error at /pytorch/aten/src/THC/THCGeneral.cpp:47
THCudaCheck FAIL file=/pytorch/aten/src/THC/THCGeneral.cpp line=47 error=3 : initialization error
Process Process-1:
Traceback (most recent call last):
File “/home/sss/anaconda3/envs/torch_new/lib/python3.8/multiprocessing/process.py”, line 315, in _bootstrap
self.run()
File “/home/sss/anaconda3/envs/torch_new/lib/python3.8/multiprocessing/process.py”, line 108, in run
self._target(*self._args, **self._kwargs)
File “/home/sss/Desktop/Experiment/test.py”, line 14, in init_process
fn(idx)
File “/home/sss/Desktop/Experiment/test.py”, line 8, in func
a = torch.randn([10, 10, 10], device=‘cuda’)
File “/home/sss/anaconda3/envs/torch_new/lib/python3.8/site-packages/torch/cuda/ **init** .py”, line 190, in _lazy_init
torch._C._cuda_init()
RuntimeError: cuda runtime error (3) : initialization error at /pytorch/aten/src/THC/THCGeneral.cpp:47

When I comment “from numba import cuda”, there is no error reported. Since I really need numba and cuda, I can’t comment them. Can anyone solve my problems?

You could either see which version of cuda numba uses see if PyTorch offers that, too (or vice versa) or you could self-compile one or the other or both.

Thank you for the reply.
I compiled a new numba for pytorch. It worked !
Thank you !