# 快速介绍

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")和torch.cuda.set_device("cuda:0")是常用的设置显卡的两种方式,二者的区别在于:

1. 使用torch.device的方式设置显卡之后,需用使用

import torch

device = torch.device("cuda:0")

model.to(device)

data = data.to(device)

 的方式将模型和数据放到GPU上

2. 如果使用torch.cuda.set_device('cuda:0')的方式,则改用

import torch

torch.cuda.set_device("cuda:0")

model.cuda()

data = data.cuda()

 的方式将模型和数据加载到显卡中

3. 即前者需要采用to()方法,而后者可以使用cuda()方法,而无需传递参数

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐