Here’s a copy of my LSTM definition
class LSTM(nn.Module):

    def __init__(self, input_dim, hidden_dim, num_layers, output_dim):

        super(LSTM, self).__init__()

        # Hidden dimensions

        self.hidden_dim = hidden_dim

        # Number of hidden layers

        self.num_layers = num_layers

        # Building your LSTM

        # batch_first=True causes input/output tensors to be of shape

        # (batch_dim, seq_dim, feature_dim)

        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True, dropout = 0.3)

        # Readout layer

        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):

        # Initialize hidden state with zeros

        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).cuda().requires_grad_()

        # Initialize cell state

        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).cuda().requires_grad_()

        # One time step

        # We need to detach as we are doing truncated backpropagation through time (BPTT)

        # If we don't, we'll backprop all the way to the start even after going through another batch

        out, (hn, cn) = self.lstm(x, (h0.detach(), c0.detach()))

        # Index hidden state of last time step

        # out.size() --> 100, 28, 100

        # out[:, -1, :] --> 100, 100 --> just want last time step hidden states! 

        out = self.fc(out[:, -1, :]) 

        # out.size() --> 100, 10

        return out

model = LSTM(input_dim=input_dim, hidden_dim=hidden_dim, output_dim=output_dim, num_layers=num_layers)

model = model.to(device)

loss_fn = torch.nn.MSELoss(size_average=True)

optimizer = torch.optim.Adam(model.parameters(), lr=0.005)


I’m using LSTM RNN to make time-series prediction based on historical steps there it uses past trajectory data to predict next position.
the problem with training as I described before is best shown in these two images below
image934×942 77.5 KB
the top graph is training on google colab T4 or our Quadro P4000
the bottom graph is training on our 3070 and 3090
the train score diverged much earlier and doesn’t perform as well
Thank you