1、如果只需要使用一块显卡
device = torch.device(‘cuda‘ if torch.cuda.is_available() else ‘cpu‘)
2、指定使用多块显卡：在python文件中指定
import os os.environ[‘CUDA_VISIBLE_DEVICES‘] = ‘0,1‘
3、在命令行中指定
CUDA_VISIBLE_DEVICES=0,1 python train.py
清空显存：
torch.cuda.empty_cache()
colab中默认是只有一块卡的。
原文：https://www.cnblogs.com/xiximayou/p/12430325.html