+

src/models_builder/gnn_models.py

@@ -835,7 +835,7 @@ def train_1_step(self, gen_dataset):
         self.gnn.eval()
         return loss.cpu().detach().numpy().tolist()
 
-    def train_on_batch(self, batch, task_type):
+    def train_on_batch(self, batch, task_type=None):
         loss = None
         if task_type == "single-graph":
             self.optimizer.zero_grad()