How to convert variable-sized 3D matrix to 2D?

Hi, everyone. My output size of TransformerEncoder is [Batch size, Sequence length, Embedding size], where the sequence length is variable. I want to convert my output size to [Batch size, 1, Embedding size].

I tried to use AvgPool1d, but kernel size is not set easily.

You could use adaptive pooling layers (nn.AdaptiveAvgPool1d), which use a variable kernel size to create the desired output shape.
Alternatively, you could also just calculate the mean via output = output.mean(1, keepdim=True).

Thank you :grinning: