Calculate alpha values with torch mean?
I'm trying to calculate the alpha values as explained here.
I have as argument a tensor with shape (1, 512, 14, 14). To calculate alpha values I need to calculate the average of all dimensions except the channel dimension, so the output will have the shape (1, k, 1, 1) which is essentialy (k,).
How can I do this in PyTorch?
Thanks!

You could permute the first and second axis to keep the channel dimension on dim=0, then flatten all other dimensions, and lastly, take the mean on that new axis:
x.permute(1, 0, 2, 3).flatten(start_dim=1).mean(dim=1)

Here are the shapes, step by step:
&gt;&gt;&gt; x.permute(1, 0, 2, 3).shape
(512, 1, 14, 14)

&gt;&gt;&gt; x.permute(1, 0, 2, 3).flatten(start_dim=1).shape
(512, 1, 196)

&gt;&gt;&gt; x.permute(1, 0, 2, 3).flatten(start_dim=1).mean(dim=1).shape
(512,)

