【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中默认是只有一块卡的。