feat: add callback support to objectives

moverseai · Oct 24, 2024 · ed12b18 · ed12b18
1 parent 37e82e8
commit ed12b18
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/moai/engine/callbacks/model.py b/moai/engine/callbacks/model.py
@@ -25,3 +25,8 @@ def __init__(
                         # get functions from operations
                         if isinstance(oper.func, Callback):
                             self.data.append(oper.func)
+            if hasattr(model, "named_objectives"):
+                for stage in model.named_objectives.values():
+                    for obj in stage.values():
+                        if isinstance(obj, Callback):
+                            self.data.append(obj)