I am using torchtext to prepare IMDB dataset . I am using the below code along with BucketIterator to generate batches.

TEXT = data.Field(lower=True, batch_first=False,fix_length=40,)
LABEL = data.Field(sequential=False,)
train, test = datasets.IMDB.splits(TEXT, LABEL)
train_iter, test_iter = data.BucketIterator.splits((train, test), batch_size=64, device=-1,shuffle=True)

The above code generates a data of the size [40,64] where 40 is the length of each document and 64 is the batch size. I am using embedding as the first layer in the network which transforms the input data to the size [40,64,300] where 300 is the dimensions. How can I use this data with an nn.LSTM layer for a text classification problem where the number of categories are 3.
Thanks,
Vishnu