【colab pytorch】指定使用的显卡

时间:2022-07-23
本文章向大家介绍【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中默认是只有一块卡的。

随机文章
本站知识点必读
最近更新