refactor scheduler

logix-project · Apr 25, 2024 · 617dc90 · 617dc90
1 parent 97e5f58
commit 617dc90
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/logix/logix.py b/logix/logix.py
@@ -81,6 +81,10 @@ def __init__(
         self.type_filter = None
         self.name_filter = None
 
+        # Deferred Imports to avoid circular imports
+        from logix.scheduler import LogIXScheduler
+        self.scheduler = LogIXScheduler(self)
+
     def watch(
         self,
         model: nn.Module,

diff --git a/logix/scheduler.py b/logix/scheduler.py
@@ -7,16 +7,17 @@ class LogIXScheduler:
     def __init__(
         self,
         logix: LogIX,
-        lora: str = "none",
-        hessian: str = "none",
-        save: str = "none",
     ):
         self.logix = logix
 
         self._epoch = -1
         self._lora_epoch = -1
         self._logix_state_schedule = []
 
+        lora = logix.config.scheduler.lora
+        hessian = logix.config.scheduler.hessian
+        save = logix.config.scheduler.save
+
         self.sanity_check(lora, hessian, save)
         self.configure_lora_epoch(lora)
         self.configure_schedule(lora, hessian, save)