Solved by ptrblck in post #9 
              
              
                Thanks for the code. 
Since you are rewrapping some modules and remove the last linear layer, the output shape will be the last activation after the average pooling layer. 
For your input shape, it’ll be [batch_size, 512, 10, 10]. 
If you want to get a single pixel in the spatial size, you would hav…