Is your feature request related to a problem? Please describe.

Goal: clean up the code making it easier for maintenance as well as trying some new architectures to see if they improve performance

Describe the solution you'd like

No response

Describe alternatives you've considered

No response

Additional context

No response

Organisation

ECMWF