From 659652f38f31c8bf49a061c2658aed26397fce0d Mon Sep 17 00:00:00 2001 From: Mario Santa Cruz Date: Fri, 18 Oct 2024 10:42:31 +0000 Subject: [PATCH] fix: typo --- src/anemoi/models/layers/graph.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/anemoi/models/layers/graph.py b/src/anemoi/models/layers/graph.py index 7a4e0d4..c3608c5 100644 --- a/src/anemoi/models/layers/graph.py +++ b/src/anemoi/models/layers/graph.py @@ -62,8 +62,8 @@ def register_fixed_attributes(self, graph_data: HeteroData) -> None: """Register fixed attributes.""" self.nodes_names = list(graph_data.node_types) self.num_nodes = {nodes_name: graph_data[nodes_name].num_nodes for nodes_name in self.nodes_names} - self.coord_dims = {2 * graph_data[nodes_name].x.shape[1] for nodes_name in self.nodes_names} - self.attr_ndims = {self.coord_dims[nodes_name] + self.num_trainable_params for nodes_name in self.nodes_names} + self.coord_dims = {nodes_name: 2 * graph_data[nodes_name].x.shape[1] for nodes_name in self.nodes_names} + self.attr_ndims = {nodes_name: self.coord_dims[nodes_name] + self.num_trainable_params for nodes_name in self.nodes_names} def register_coordinates(self, name: str, node_coords: torch.Tensor) -> None: """Register coordinates."""