【colab pytorch】指定使用的顯卡

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中默認是只有一塊卡的。