【colab pytorch】指定使用的显卡

西西嘛呦

发布于 2020-08-26 10:15:32

1.4K00

代码可运行

运行总次数：0

代码可运行

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

本文参与腾讯云自媒体同步曝光计划，分享自作者个人站点/博客。

原始发表：2020-03-06 ，如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自作者个人站点/博客前往查看

如有侵权，请联系 cloudcommunity@tencent.com 删除。

本文参与腾讯云自媒体同步曝光计划，欢迎热爱写作的你一起参与！

登录后参与评论

0 条评论

热度