Add ParamsT import

Signed-off-by: Alexandros Koumparoulis <[email protected]>
NVIDIA · Nov 5, 2024 · 030076e · 030076e
1 parent 8732e2a
commit 030076e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/lightning/pytorch/optim/pytorch.py b/nemo/lightning/pytorch/optim/pytorch.py
@@ -20,7 +20,7 @@
 
 from nemo.lightning.megatron_parallel import MegatronParallel
 from nemo.lightning.pytorch.optim.base import LRSchedulerModule, OptimizerModule
-
+from torch.optim.optimizer import ParamsT
 
 def _param_does_not_have_wd(param_name, param):
     return 'bias' in param_name