关于pytorch1.9.0中cpp_extension的默认CUDA版本是11.0的问题

今天在用CUDAExtension编译一个算子的时候,发现它默认的CUDA路径指向的是11.0版本,而我的CUDA版本是11.1.

这会导致使用nvcc编译时,会报错No such file or directory错误。

这时我们只需要将CUDA的路径指向我们安装的路径即可,比如下面这样

torch.utils.cpp_extension.CUDA_HOME = "/usr/local/cuda-11.1"

就能够将CUDA路径纠正。

猜你喜欢

转载自blog.csdn.net/qq_19313495/article/details/121765916