pytorch - RuntimeError: CUDA error: invalid argument torch tensor.cuda() 在执行编译的 C++ cuda 函数后引发错误

Question

我所有的运营商都是：

1. .cu在文件中写一个函数。
2. 编译.cu使用nvcc xxx.cu -o xxx.so -shared -Xcompiler -fPIC
3. 加载它ctypes.cdll

错误发生在：

lib = ctypes.cdll.load('xxx.so')

tmp = torch.ones((1, 1)).cuda()

a = np.ones((1, 1))
b = lib.xxx(a)

tmp = torch.ones((1, 1)).cuda()

执行第二个时出错tensor.cuda()

RuntimeError: CUDA error: invalid argument

但即使发生错误，torch.cuda.is_available()仍然是 True。

如果我从不将张量加载到 gpu，错误就会消失，即：

lib = ctypes.cdll.load('xxx.so')

tmp = torch.ones((1, 1))

a = np.ones((1, 1))
b = lib.xxx(a)

tmp = torch.ones((1, 1)).cuda()

运行没有错误。

我想找到一种在.so任何地方执行该功能的方法。

标签： pytorchshared-libraries