return epoch in scheduler

logix-project · Nov 23, 2023 · bd27357 · bd27357
1 parent 2fe6a65
commit bd27357
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/analog/scheduler.py b/analog/scheduler.py
@@ -83,7 +83,7 @@ def __next__(self):
                 self.analog.ekfac()
             if self._epoch == self.execution_schedule["lora"]:
                 self.analog.add_lora()
-            return
+            return self._epoch
         raise StopIteration
 
     def __len__(self):