Index of the last timestep LSTM

1

Hi! I’m working with LSTM and the output tensor has a shape of (batch_size, seq_length, hidden_size). To get only the last output of the network, I need to get the last index of the tensor output[:, seq_length -1, :] or the first one output[:, 0, :] ?

Thanks!

2

You can get the last time step via output[:. -1] assuming you’ve used batchj_first=True during the setup of the nn.LSTM.

1 Like