Pytorch的trick与技巧（坑）

2024年3月27日 170次阅读来源: 好奇的小will

版本：pytorch 1.0

指定GPU显卡

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

设定全局的tensor 初始化为 tensor.cuda()

if torch.cuda.is_available():
    torch.set_default_tensor_type('torch.cuda.FloatTensor')

将CPU转为GPU的tensor

# self.model = TextCNNNet(self.kwargs).cuda()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    原文作者：好奇的小will
    原文地址: https://zhuanlan.zhihu.com/p/59216652
    本文转自网络文章，转载此文章仅为分享知识，如有侵权，请联系博主进行删除。