chore: add frequency parameter to wrap_lr_scheduler [DET-5148] (#2087)

determined-ai · Mar 22, 2021 · 3d2f54f · 3d2f54f
1 parent 18c3994
commit 3d2f54f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/harness/determined/pytorch/_pytorch_context.py b/harness/determined/pytorch/_pytorch_context.py
@@ -211,6 +211,7 @@ def wrap_lr_scheduler(
         self,
         lr_scheduler: torch.optim.lr_scheduler._LRScheduler,
         step_mode: pytorch.LRScheduler.StepMode,
+        frequency: int = 1,
     ) -> torch.optim.lr_scheduler._LRScheduler:
         """
         Returns a wrapped LR scheduler.
@@ -244,7 +245,7 @@ def wrap_lr_scheduler(
                 self.optimizers,
                 "Must use an optimizer that is returned by wrap_optimizer()",
             )
-        wrapped = pytorch.LRScheduler(lr_scheduler, step_mode)
+        wrapped = pytorch.LRScheduler(lr_scheduler, step_mode, frequency)
         self.lr_schedulers.append(wrapped)
 
         # Return the original LR scheduler to the user in case they have customizations that we