add missing kwarg to BaseHiGraphModel.__init__

mllam · Sep 27, 2024 · 5df2ecf · 5df2ecf
1 parent b7bf506
commit 5df2ecf
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/neural_lam/models/base_hi_graph_model.py b/neural_lam/models/base_hi_graph_model.py
@@ -12,8 +12,8 @@ class BaseHiGraphModel(BaseGraphModel):
     Base class for hierarchical graph models.
     """
 
-    def __init__(self, args):
-        super().__init__(args)
+    def __init__(self, args, datastore):
+        super().__init__(args, datastore=datastore)
 
         # Track number of nodes, edges on each level
         # Flatten lists for efficient embedding