GPU Functions Are Broken

[seanlaw]

Sadly, numba.cuda no longer works and numba-cuda has gone into maintenance mode and is being replaced by numa-cuda-mlir:

numba/numba#10753
nzxyin/sdtw-cuda-torch#2

The solution would be to do:

try:
    from numba_cuda_mlir import cuda
except ModuleNotFoundError:
    from numba import cuda