PyTorch torch.cuda.device_count 返回值与实际 GPU 数量不一致

今天执行基于 PyTorch 的图像分类算法程序时，触发了自己写的断言错误。而断言的细节，就是判断用户输入的 GPU 编号是否合法。
调试打开，发现 torch.cuda.device_count() 返回的是 1。而我机器上明明是两张卡。

一脸懵逼。
查阅 PyTorch 官网后，发现是使用问题。我在调用 device_count 之前，已经设置过了环境变量 CUDA_VISIBLE_DEVICES。
通过在 os.environ["CUDA_VISIBLE_DEVICES"] 代码之前执行 device_count, 发现返回的是 2。至此，问题已定位。
PS. 官方推进使用 os.environ["CUDA_VISIBLE_DEVICES"] 的形式来设定使用的 GPU 显卡。

参考链接

• I have 8 GPUs ,buyt why my torch.cuda.device_count() return only 1?