Numbapro or Numba for PyTorch Cuda Extension

Can I use the numbapro or some other python cuda library to write the cuda python code for the pytorch extension?
I find that the present way to write pytorch extension code running on gpu is mostly writing the cuda file with C language. I want to directly write the cuda file with python in pytorch.

1 Like

not numba/numbapro, but you can use cupy: