If you just want the shape to (1x20), you could just do return output.mean(dim=0). It should be noted that taking the mean with make your network dependent on the number of samples!
But shouldn’t the network take in N samples and return N outputs? Then your loss function with sum over the number of samples and return a single value?



 Mole_Turner:

At present this isn’t happening and it’s outputting 5 x 20


The reason why your network returns 5x20 as the size is due to matrix multiplication. For example with A times B the shapes must be ‘compatible’. So, the number of number of columns in A must match the number of rows in B, (the inner two dimensions in your example, [5,2] x [2,128]), and the remaining dimensions give the output shape from the multiplication (which is a matrix of shape [5,128]).