Average Pool to get rid of max_len dim

I have a tensor with shape (batch_size, max_len, latent_dim) I want to make this tensor of shape (batch_size, latent_dim). I want to get rid of max_len dim.

Why I want to do this?

I am doing sentiment analysis so initially, I have tensor of shape (batch_size, max_len) → then I gave this to embedding so it became (batch_size, max_len, emb_dim) → then I gave it to LSTM so it became (batch_size, max_len, latent_dim) → then I gave it to the final dense layer with 2 neurons for ‘positive’ and ‘negative’ sentiment so I can have two classes but I am getting shape (batch_size, max_len, 2) but I want (batch_size, 2) so then I can come to a conclusion what class my neural network has predicted.

In the tensorflow tutorial here you can see they have done some average pooling so (batch_size, max_len, latent_dim) becomes (batch_size, latent_dim). But I want to use Pytorch, how can I achieve this?

Thanks!