Pytorch检查cuda

1. 是否可以用显卡：

torch.cuda.is_available()

2. 显卡个数：

torch.cuda.device_count()

3. 显卡名字：

torch.cuda.get_device_name(0)

 

设置可见GPU:

CUDA_VISIBLE_DEVICES=1 python run.py

或者：

import os
os.environ["CUDA_VISIBLE_DEVICES"]="1"