Update Lightning import (#2549)

* update import Signed-off-by: Maanu Grover <[email protected]> * run formatting --------- Signed-off-by: Maanu Grover <[email protected]>
EleutherAI · Dec 9, 2024 · 0b99443 · 0b99443
1 parent 2d11f2e
commit 0b99443
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/lm_eval/models/nemo_lm.py b/lm_eval/models/nemo_lm.py
@@ -187,11 +187,11 @@ def __init__(
         **kwargs,
     ):
         try:
+            from lightning.pytorch.trainer.trainer import Trainer
             from nemo.collections.nlp.modules.common.text_generation_utils import (
                 generate,
             )
             from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy
-            from pytorch_lightning.trainer.trainer import Trainer
 
             self.generate = generate
         except ModuleNotFoundError as exception: