Since your target batch size us 128, and your initial batch size is 128 according to one of your comments
# inputs shape [128, 256] = [batch size, max seq len]

somewhere you make a mistake with the dimensions. My first guess would be that you define your nn.LSTM with ‘batch_first=False’ (default value). If you give the LSTM layer a shape of [128, 256], 128 is the sequence length and 256 the batch size. So you either have to define your LSTM with batch_first=True or add
inputx = inputx.transpose(1,0)

before giving it to the LSTM. Of course, these change might affect subsequent layers which you then also have to modify.