Question about .mean() function

There is a WGAN implementation in Pytorch(the official one). https://github.com/martinarjovsky/WassersteinGAN/blob/master/models/dcgan.py
I am unsure whats happening around line 52, where we are supposed to do a global average pooling.
The last layer is :
main.add_module('final.{0}-{1}.conv'.format(cndf, 1),nn.Conv2d(cndf, 1, 4, 1, 0, bias=False))

and in the forward function, there is :

output = output.mean(0)
return output.view(1)```

what does output = output.mean(0) do ? 

I read the docs and mean returns average of supplied dimension. So i created a similar shaped tensor as I would expect from the last convolutional layer in the model above. 
The shape was (1x1x10x10 ).  batch size * one layer output * width* height

when i run .mean(0) on this tensor, it does NOT return the global mean of the 10x10 matrix. 

>     In [38]: a
>     Out[38]:

>     (0 ,0 ,.,.) =
>       0.3912  0.8033  0.3859
>       0.0037  0.1687  0.2818
>       0.2725  0.4355  0.2085
>     [torch.FloatTensor of size 1x1x3x3]

>     In [39]: a.mean(0)
>     Out[39]:

>     (0 ,0 ,.,.) =
>       0.3912  0.8033  0.3859
>       0.0037  0.1687  0.2818
>       0.2725  0.4355  0.2085
>     [torch.FloatTensor of size 1x1x3x3]

I can't understand whats wrong. Help is greatly appreciated

If you have tensor A with shape of [1, a, b, c] then calling A.mean(0) will return you new tensor with averaged over zero-axis elements.
Since you have A.size(0) = 1 nothing will happen after doing so.

If you want to average whole tensor you have to do it by A.mean()

9 Likes