if 51 is time dimension, you need to use nn.LSTM(256, 256, batch_first=True)
first output of lstm module contains all time steps, for tasks like yours you can discard non-final outputs by slicing (out1 = out1[:,50]), or use one of tensors from second output (tuple named ‘hidden’ in your code), as these contain non “unrolled” final states too. In your simple case (one rnn layer, same sequence lengths) these methods are equivalent.