您的当前位置：首页 torch.nn.module.to()中参数为i与参数为‘cuda:i‘相同

torch.nn.module.to()中参数为i与参数为‘cuda:i‘相同

来源：爱问旅游网

代码如下：

import torch.nn as nn

class Model(nn.Module):
    def __init__(self, input_size, output_size):
        super(Model, self).__init__()  # 继承父类的方法
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.fc(input)
        return output

import torch
model = Model(10, 2)
print(next(model.parameters()).device)

结果：cpu

device = torch.device('cuda:1')
model = model.to(device)
print(next(model.parameters()).device)

结果：cuda:1

model = model.to('cuda:2')
print(next(model.parameters()).device)

结果：cuda:2

model = model.to(3)
print(next(model.parameters()).device)

结果：cuda:3
model.to(1)和model.to('cuda:1')效果一致。

因篇幅问题不能全部显示，请点此查看更多更全内容

查看全文