add HeteroData import

ecmwf · Feb 5, 2025 · 6fbfe4c · 6fbfe4c
1 parent 11d8222
commit 6fbfe4c
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/graphs/tests/edges/test_edge_attributes.py b/graphs/tests/edges/test_edge_attributes.py
@@ -10,12 +10,16 @@
 
 import pytest
 import torch
+from typing import TYPE_CHECKING
 
 from anemoi.graphs.edges.attributes import AttributeFromSourceNode
 from anemoi.graphs.edges.attributes import AttributeFromTargetNode
 from anemoi.graphs.edges.attributes import EdgeDirection
 from anemoi.graphs.edges.attributes import EdgeLength
 
+if TYPE_CHECKING:
+    from torch_geometric.data import HeteroData
+
 TEST_EDGES = ("test_nodes", "to", "test_nodes")