I got better performance by replacing the ED_expansion and decoder_inv modules with ResidualBlock.
I'd like to double check that this preserves the resolution invariance of CNO, since now there are no Invariant blocks. Removing ED_expansion makes it so that the input and output resolution of the model must match, but I'm ok with that as long as the model is still a neural operator, i.e., it can handle resolutions different from the training resolution (as long as the bandlimit isn't reduced by reducing the resolution too much) without changing the performance.