pytorch调用gpu

第一步！指定gpu

import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0'

第二步！

对于每一个要踹到gpu去的Tensor或者model x

使用x = x.cuda()就ok了

嘤嘤嘤

posted @ 2020-04-21 12:06  地球上最后一个直男  阅读(1013)  评论(0)  编辑  收藏  举报