diff --git a/src/power_grid_model_ds/_core/model/graphs/models/base.py b/src/power_grid_model_ds/_core/model/graphs/models/base.py
index e717b4a..8e7e174 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/base.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/base.py
@@ -6,7 +6,6 @@
 from contextlib import contextmanager
 from typing import Generator
 
-import numpy as np
 from numpy._typing import NDArray
 
 from power_grid_model_ds._core.model.arrays.pgm_arrays import Branch3Array, BranchArray, NodeArray
@@ -183,7 +182,7 @@ def delete_branch3_array(self, branch_array: Branch3Array, raise_on_fail: bool =
             self.delete_branch_array(branches, raise_on_fail=raise_on_fail)
 
     @contextmanager
-    def tmp_remove_nodes(self, nodes: list[int] | NDArray[np.int32]) -> Generator:
+    def tmp_remove_nodes(self, nodes: list[int]) -> Generator:
         """Context manager that temporarily removes nodes and their branches from the graph.
         Example:
             >>> with graph.tmp_remove_nodes([1, 2, 3]):
@@ -248,7 +247,7 @@ def get_all_paths(self, ext_start_node_id: int, ext_end_node_id: int) -> list[li
 
         return [self._internals_to_externals(path) for path in internal_paths]
 
-    def get_components(self, substation_nodes: NDArray[np.int32]) -> list[list[int]]:
+    def get_components(self, substation_nodes: list[int]) -> list[list[int]]:
         """Returns all separate components when the substation_nodes are removed of the graph as lists"""
         with self.tmp_remove_nodes(substation_nodes):
             internal_components = self._get_components()