fix magnitude issue

GFNOrg · Feb 19, 2025 · a6ffc72 · a6ffc72
1 parent af26f6e
commit a6ffc72
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tutorials/examples/train_graph_ring.py b/tutorials/examples/train_graph_ring.py
@@ -354,6 +354,7 @@ def forward(self, states_tensor: TensorDict) -> torch.Tensor:
 
         # This is n_nodes ** 2, for each graph.
         edge_index = torch.einsum("bnf,bmf->bnm", edge_feature, edge_feature)
+        edge_index = edge_index / torch.sqrt(torch.tensor(self.hidden_dim))
 
         # Undirected.
         if self.is_directed: