fix

Signed-off-by: Alexandros Koumparoulis <[email protected]>
NVIDIA · Nov 5, 2024 · 9752704 · 9752704
1 parent 030076e
commit 9752704
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/nemo/collections/llm/recipes/optim/adam.py b/nemo/collections/llm/recipes/optim/adam.py
@@ -20,7 +20,6 @@
 from nemo.lightning.pytorch.optim import (
     CosineAnnealingScheduler,
     MegatronOptimizerModule,
-    OptimizerModule,
     PytorchOptimizerModule,
 )
 
@@ -35,7 +34,7 @@ def distributed_fused_adam_with_cosine_annealing(
     max_lr: float = 1e-4,
     min_lr: Optional[float] = None,
     clip_grad: float = 1.0,
-) -> run.Config[OptimizerModule]:
+) -> run.Config[PytorchOptimizerModule]:
 
     opt_cfg = run.Config(
         OptimizerConfig,
@@ -72,7 +71,7 @@ def pytorch_adam_with_cosine_annealing(
     constant_steps: int = 0,
     max_lr: float = 1e-5,
     min_lr: Optional[float] = None,
-) -> run.Config[OptimizerModule]:
+) -> run.Config[PytorchOptimizerModule]:
     from torch.optim import Adam
 
     return run.Config(
@@ -96,7 +95,7 @@ def pytorch_adam_with_cosine_annealing(
 @run.cli.factory
 def pytorch_adam_with_flat_lr(
     lr: float = 1e-5,
-) -> run.Config[OptimizerModule]:
+) -> run.Config[PytorchOptimizerModule]:
     from torch.optim import Adam
 
     return run.Config(

diff --git a/nemo/collections/llm/recipes/optim/sgd.py b/nemo/collections/llm/recipes/optim/sgd.py
@@ -16,7 +16,7 @@
 
 import nemo_run as run
 
-from nemo.lightning.pytorch.optim import CosineAnnealingScheduler, OptimizerModule, PytorchOptimizerModule
+from nemo.lightning.pytorch.optim import CosineAnnealingScheduler, PytorchOptimizerModule
 
 
 @run.cli.factory
@@ -26,7 +26,7 @@ def pytorch_sgd_with_cosine_annealing(
     max_lr: float = 1e-5,
     min_lr: Optional[float] = None,
     wd: float = 1e-4,
-) -> run.Config[OptimizerModule]:
+) -> run.Config[PytorchOptimizerModule]:
     from torch.optim import SGD
 
     return run.Config(
@@ -49,7 +49,7 @@ def pytorch_sgd_with_cosine_annealing(
 def pytorch_sgd_with_flat_lr(
     lr: float = 1e-5,
     wd: float = 1e-4,
-) -> run.Config[OptimizerModule]:
+) -> run.Config[PytorchOptimizerModule]:
     from torch.optim import SGD
 
     return run.Config(