PytorchのLSTMでのGPU使用方法について

python
1device = torch.device("cuda:0")
2cpu = torch.device("cpu")
3・
4・
5・
6#モデル構築
7class LSTM(nn.Module):
8    def __init__(self, input_size=1, hidden_layer_size=50, output_size=1):
9        super().__init__()
10        
11        self.hidden_layer_size = hidden_layer_size
12        self.lstm = nn.LSTM(input_size, hidden_layer_size)
13        self.linear = nn.Linear(hidden_layer_size, output_size)
14        
15        self.hidden_cell = (torch.zeros(1, 1, self.hidden_layer_size).to(device),
16                            torch.zeros(1, 1, self.hidden_layer_size).to(device))
17
18    def forward(self, input_seq):
19        lstm_out, self.hidden_cell = self.lstm(input_seq.reshape(len(input_seq), 1, 1),
20                                               self.hidden_cell)
21        pred = self.linear(lstm_out.reshape(len(input_seq), -1))
22        
23        return pred[-1]
24
25torch.manual_seed(123)
26model = LSTM().to(device)
27criterion = nn.MSELoss()
28optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
29
30for epoch in range(epochs):
31    for sequence_in , t in train_dataset:
32        sequence_in = sequence_in.to(device)
33        t = t.to(device)
34        y = model(sequence_in)
35        y = y.to(device)
36        loss = criterion(y, t)
37        
38        optimizer.zero_grad()
39        model.hidden_cell = (torch.zeros(1, 1, model.hidden_layer_size).to(device),
40                             torch.zeros(1, 1, model.hidden_layer_size).to(device))
41        
42        loss.backward()
43        optimizer.step()
44        
45    if epoch % 10 == 0:
46        loss_list.append(loss.detach().to(cpu))
47        print(f'Epoch {epoch+1} Loss {loss.item()}')