From 87d90a861dd6c3440f2026f0331f96a9762b6e51 Mon Sep 17 00:00:00 2001
From: Andreas Michael Hermansen
 <97125645+AMHermansen@users.noreply.github.com>
Date: Sat, 9 Sep 2023 21:23:28 +0200
Subject: [PATCH 1/2] Change type hint from factory method to class

Changed the type hint, to refer to the class and not a factory method.
---
 src/graphnet/models/graphs/nodes/nodes.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/graphnet/models/graphs/nodes/nodes.py b/src/graphnet/models/graphs/nodes/nodes.py
index b98ad375a..6b3443e0c 100644
--- a/src/graphnet/models/graphs/nodes/nodes.py
+++ b/src/graphnet/models/graphs/nodes/nodes.py
@@ -68,5 +68,5 @@ def _construct_nodes(self, x: torch.tensor) -> Data:
 class NodesAsPulses(NodeDefinition):
     """Represent each measured pulse of Cherenkov Radiation as a node."""
 
-    def _construct_nodes(self, x: torch.tensor) -> Data:
+    def _construct_nodes(self, x: torch.Tensor) -> Data:
         return Data(x=x)

From f6ee8c107e5b10f5cfe949a74da177190bc73c7b Mon Sep 17 00:00:00 2001
From: Andreas Michael Hermansen
 <97125645+AMHermansen@users.noreply.github.com>
Date: Sat, 9 Sep 2023 21:25:46 +0200
Subject: [PATCH 2/2] Fix type hint to refer to class name and not factory
 method

---
 src/graphnet/models/graphs/graph_definition.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/graphnet/models/graphs/graph_definition.py b/src/graphnet/models/graphs/graph_definition.py
index dfb2cedac..c29867155 100644
--- a/src/graphnet/models/graphs/graph_definition.py
+++ b/src/graphnet/models/graphs/graph_definition.py
@@ -71,7 +71,7 @@ def __init__(
 
     def forward(  # type: ignore
         self,
-        node_features: np.array,
+        node_features: np.ndarray,
         node_feature_names: List[str],
         truth_dicts: Optional[List[Dict[str, Any]]] = None,
         custom_label_functions: Optional[Dict[str, Callable[..., Any]]] = None,