烦人的CUDA error：-UserWarning: CUDA initialization: CUDA unknown error

参考：pytorch学习笔记-UserWarning: CUDA initialization: CUDA unknown error - 简书

跑yolov5时，安装好pytorch之后，运行脚本

import torch
a=torch.cuda.is_available()
print(torch.__version__)
print(a)

出现了错误：

UserWarning: CUDA initialization: CUDA unknown error - this may be due to an incorrectly set up environment, 
e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the available devices to be zero.
 (Triggered internally at  /opt/conda/conda-bld/pytorch_1623448255797/work/c10/cuda/CUDAFunctions.cpp:115.)
  return torch._C._cuda_getDeviceCount() > 0

首先检查显卡驱动，CUDA，cudnn，以及pytorch的版本是否匹配，如果不匹配，需要卸载之后重装对应的版本。

CUDA10.2 Python3.8 pytorch1.8 没有错。

如果版本都是正确的，需要设置一下环境变量，输入sudo vim ~/.bashrc，在最后加上：

# 前三行是安装CUDA的时候需要设置的
export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
export CUDA_HOME=/usr/local/cuda/bin

export CUDA_VISIBLE_DEVICES=0

保存退出，试一下是否能够使用CUDA。

如果还是不行，输入apt-get install nvidia-modprobe，应该就没有问题了。

我是在安装了apt-get install nvidia-modprobe之后才OK的。