Pytorch - Apply pooling on specific dimension
I have a 3 dimension vector. I would like to perform a 1d max pool on the second dimension.

According to the documentation of pytorch the pooling is always performed on the last dimension.


  https://pytorch.org/docs/stable/nn.html#maxpool1d


For example:

&gt;&gt;&gt; x = torch.rand(5, 64, 32)
&gt;&gt;&gt; pool = nn.MaxPool1d(2, 2)
&gt;&gt;&gt; pool(x).shape
torch.Size([5, 64, 16])


My desired output:

torch.Size([5, 32, 32])


How can i do that?

You can simply permute the dimensions:

x = torch.rand(5, 128, 32)
pool = nn.MaxPool1d(2, 2)
pool(x.permute(0,2,1)).permute(0,2,1)  # shape (5, 128, 32) -&gt; (5, 64, 32)

