return torch._C._cuda_getDeviceCount() ＞ 0 卡住了

今天遇到一个问题，在命令行模式下通过ssh远程调试pytorch代码，后面就卡主了，无论重启远程实例还是IDE都不行。

1、通过以下命令追踪代码：

python -m trace --trace YOURSCRIPT.py

发现是执行 return torch._C._cuda_getDeviceCount() ＞ 0 就卡主了。

ps -ef |grep python，发现调试进程还在

nvidia-smi，发现并没有进程

2、网上查资料找到如下答案

参考：CUDA fails to reinitialize after system suspend - PyTorch Forums

It’s a known issue of Linux’ suspend mode and the CUDA driver.
Resetting the device as described here 446 might work.

You could try to execute:

sudo rmmod nvidia_uvm
sudo modprobe nvidia_uvm

which helps on my Ubuntu system after it was suspended.

意思是说，这是linux的一个已知问题

3、如何安装rmmod和modprobe命令

apt-get install kmod

4、执行命令，报

rmmod: ERROR: Module nvidia_uvm is in use

5、查看有那些进程在使用：sudo lsof -n -w /dev/nvidia*

kill这些进程