@ptrblck even after pooling wouldn’t the size be [batch, classes,1] whereas the loss function would be expecting [batch, classes]
I’m not able to understand how to remove that last dimension altogether