Improve prediction heads #529
Description
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