type hints

jla-gardner · Feb 12, 2024 · 82453c5 · 82453c5
1 parent 3c1fc3a
commit 82453c5
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/src/graph_pes/transform.py b/src/graph_pes/transform.py
@@ -58,6 +58,14 @@ def forward(
             The transformed data.
         """
 
+    # add type hints to play nicely with mypy
+    def __call__(
+        self,
+        x: Shaped[Tensor, "shape ..."],
+        graph: AtomicGraph | AtomicGraphBatch,
+    ) -> Shaped[Tensor, "shape ..."]:
+        return super().__call__(x, graph)
+
     @abstractmethod
     def inverse(
         self,