I’m trying to train a LSTM connected to couple MLP layers. The model is coded as follows:
class RNNBlock(nn.Module):

def __init__(self, in_dim, hidden_dim, num_layer=1, dropout=0):
    super(RNNBlock, self).__init__()

    self.hidden_dim = hidden_dim
    self.num_layer = num_layer
    
    self.lstm = nn.LSTM(in_dim, hidden_dim, num_layer, dropout)

def forward(self, onehot, length):
    batch_size = onehot.shape[0]
    
    h_in = torch.randn(self.num_layer, batch_size, self.hidden_dim).cuda()
    c_in = torch.randn(self.num_layer, batch_size, self.hidden_dim).cuda()
    packed = nn.utils.rnn.pack_padded_sequence(onehot, length, batch_first=True).cuda()
    
    output, (h_out, c_out) = self.lstm(packed, (h_in, c_in))
    unpacked, unpacked_length = nn.utils.rnn.pad_packed_sequence(output, batch_first=True)
    vectors = list()
    for i, vector in enumerate(unpacked):
        vectors.append(unpacked[i, unpacked_length[i]-1, :].view(1, -1))
    out = torch.cat(vectors, 0)
    return out

class Predictor(nn.Module):

def __init__(self, in_dim, out_dim, act=None):
    super(Predictor, self).__init__()

    self.linear = nn.Linear(in_dim, out_dim)
    nn.init.xavier_normal_(self.linear.weight)
    self.activation = act

def forward(self, x):
    out = self.linear(x)
    if self.activation != None:
        out = self.activation(out)
    return out

class RNNNet(nn.Module):

def __init__(self, args):
    super(RNNNet, self).__init__()

    self.rnnBlock = RNNBlock(args.in_dim, args.hidden_dim, args.num_layer, args.dropout)
    self.pred1 = Predictor(args.hidden_dim, args.pred_dim1, act=nn.ReLU())
    self.pred2 = Predictor(args.pred_dim1, args.pred_dim2, act=nn.ReLU())
    self.pred3 = Predictor(args.pred_dim2, args.out_dim)

def forward(self, onehot, length):
    out = self.rnnBlock(onehot, length)
    out = self.pred1(out)
    out = self.pred2(out)
    out = self.pred3(out)
    return out

and this is my train function
def train(model, device, optimizer, criterion, data_train, bar, args):
epoch_train_loss = 0
epoch_train_mae = 0

for i, batch in enumerate(data_train):
    list_onehot = torch.tensor(batch[0]).cuda().float()
    list_length = torch.tensor(batch[1]).cuda()
    list_logP = torch.tensor(batch[2]).cuda().float()
    # Sort onehot tensor with respect to the sequence length.
    list_length, list_index = torch.sort(list_length, descending=True)
    list_length.cuda()
    list_index.cuda()
    list_onehot = torch.Tensor([list_onehot.tolist()[i] for i in list_index]).cuda().float()

    model.train()
    optimizer.zero_grad()

    list_pred_logP = model(list_onehot, list_length).squeeze().cuda()
    list_pred_logP.require_grad = False

    train_loss = criterion(list_pred_logP, list_logP)
    train_mae = mean_absolute_error(list_pred_logP.tolist(), list_logP.tolist())
    epoch_train_loss += train_loss.item()
    epoch_train_mae += train_mae

    train_loss.backward()
    optimizer.step()

    bar.update(len(list_onehot))

epoch_train_loss /= len(data_train)
epoch_train_mae /= len(data_train)

return model, epoch_train_loss, epoch_train_mae

The list_onehot and list_length tensors are loaded from the DataLoader and uploaded to GPU. Then, to use packed sequence as input, I’ve sorted the both list_onehot and list_length and uploaded to GPU. The model was uploaded to GPU and h_in, c_in tensors and packed sequence object were also uploaded to the GPU. However, when I try to run this code, it does not use GPU but only use CPU. What should I do to use GPU to train this model?