关于torch.to(device)是否赋值的坑

在我们用GPU跑程序时，需要在程序中把变量和模型放到GPU里面。
有一些坑需要注意，本博文用RNN模型实例
首先，定义device

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

对于变量，需要进行赋值操作才能真正转到GPU上：

all_input_batch=all_input_batch.to(device)

对于模型，不需要进行赋值：

 model = TextRNN()
 model.to(device)

对模型进行to(device)，还有一种方法，就是在定义模型的时候全部对模型网络参数to(device)，这样就可以不需要model.to(device)这句话。

class TextRNN(nn.Module):

    def __init__(self):
        super(TextRNN, self).__init__()
        #self.cnt = 0
        self.C = nn.Embedding(n_class, embedding_dim=emb_size,device=device)
        self.rnn = nn.RNN(input_size=emb_size, hidden_size=n_hidden,device=device)
        self.W = nn.Linear(n_hidden, n_class, bias=False,device=device)
        self.b = nn.Parameter(torch.ones([n_class])).to(device)


    def forward(self, X):
        X = self.C(X)
        #print(X.is_cuda)
        X = X.transpose(0, 1) # X : [n_step, batch_size, embeding size]
        outputs, hidden = self.rnn(X)
        # outputs : [n_step, batch_size, num_directions(=1) * n_hidden]
        # hidden : [num_layers(=1) * num_directions(=1), batch_size, n_hidden]
        outputs = outputs[-1] # [batch_size, num_directions(=1) * n_hidden]
        model = self.W(outputs) + self.b # model : [batch_size, n_class]
        return model

完整的代码太多了，不方便放在博客上
有问题或者建议可以评论