What happened?

The following lines in anemoi/training/train/forecaster.py, break the training functionality if a hierarchical graph is used:

# Check if the model is a stretched grid
if graph_data["hidden"].node_type == "StretchedTriNodes":
       mask_name = config.graph.nodes.hidden.node_builder.mask_attr_name
       limited_area_mask = graph_data[config.graph.data][mask_name].squeeze().bool()
else:
       limited_area_mask = torch.ones((1,))

Hierarchical graphs have a graph structure defined as:

data: "data"
hidden:
 - "hidden_1"
 - "hidden_2"
 - "hidden_3"

therefore cannot access graph_data["hidden"], as it does not exist.
In case of hierarchical graphs multiple masks can be computed, each per dept.

What are the steps to reproduce the bug?

Run any training config using a hierarchical graph architecture.

Version

training-0.3.3

Platform (OS and architecture)

Linux

Relevant log output

Accompanying data

No response

Organisation

MeteoSwiss