Code QuizIntermediate Selecting the LSTM Last Timestep
Spot the indexing bug when grabbing the final timestep from a batch-first LSTM output.
Codepython
import torch
import torch.nn as nn
lstm = nn.LSTM(input_size=10, hidden_size=20, batch_first=True)
# x shape: (batch=32, seq_len=5, features=10)
x = torch.randn(32, 5, 10)
output, (h_n, c_n) = lstm(x)
# output shape: (batch=32, seq_len=5, hidden=20)
# grab the last timestep for classification
last = output[-1, :, :]
print(last.shape)
What is the bug in this code that selects the last timestep?