diff --git a/aiida/tools/visualization/graph.py b/aiida/tools/visualization/graph.py
index b864ad28bb..b11ce15ff7 100644
--- a/aiida/tools/visualization/graph.py
+++ b/aiida/tools/visualization/graph.py
@@ -428,16 +428,17 @@ def edges(self):
         """return a copy of the edges"""
         return self._edges.copy()
 
-    @staticmethod
-    def _load_node(node):
+    def _load_node(self, node):
         """ load a node (if not already loaded)
 
         :param node: node or node pk/uuid
         :type node: int or str or aiida.orm.nodes.node.Node
         :returns: aiida.orm.nodes.node.Node
         """
-        if isinstance(node, (int, str)):
-            return orm.load_node(node)
+        if isinstance(node, int):
+            return orm.Node.Collection.get_cached(orm.Node, self._backend).get(pk=node)
+        if isinstance(node, str):
+            return orm.Node.Collection.get_cached(orm.Node, self._backend).get(uuid=node)
         return node
 
     @staticmethod