LayerNorm after last skip connection

[SavvasMel]

Description

This PR changes the MLP of the last block of the forecast engine. Specifically, it includes a LayerNorm with scale and bias turned off after the last skip connection of the last block.

Issue Number

This is a draft PR

Checklist before asking for review

• I have performed a self-review of my code
• My changes comply with basic sanity checks:
  • I have fixed formatting issues with ./scripts/actions.sh lint
  • I have run unit tests with ./scripts/actions.sh unit-test
  • I have documented my code and I have updated the docstrings.
  • I have added unit tests, if relevant
• I have tried my changes with data and code:
  • I have run the integration tests with ./scripts/actions.sh integration-test
  • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
  • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
• I have informed and aligned with people impacted by my change:
  • for config changes: the MatterMost channels and/or a design doc
  • for changes of dependencies: the MatterMost software development channel

[MatKbauer]

Thanks for sharing your code, Savvas. Let's please modify it slightly to minimize redundant lines. See my suggestions below.

[MatKbauer]

For the sake of having less redundant code, can we modify the MLP to receive an additional argument with_residual_layer_norm=False instead of introducing FEMLP? When calling MLP here, we can set with_residual_layer_norm=(i + 1) == self.cf.fe_num_blocks to add the residual layer norm in the last MLP layer.

[clessig]

Why would we need to modify the MLP? It can--and in my opinion definitely should--be implemented in the forecast engine. Where we can just have the LayerNorm as the last block.

[SavvasMel]

Can you check now? Is that better?