bot happy

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
NVIDIA · Sep 27, 2024 · 706e9c1 · 706e9c1
1 parent c3699c7
commit 706e9c1
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/tests/collections/llm/megatron_mixtral_pretraining.py b/tests/collections/llm/megatron_mixtral_pretraining.py
@@ -1,12 +1,10 @@
 import argparse
-import sys
 import os
 import torch
 
 from pathlib import Path
-import torch.nn.functional as F
 
-from nemo.lightning import Trainer, MegatronStrategy, MegatronMixedPrecision
+from nemo.lightning import Trainer, MegatronStrategy
 from nemo.collections.llm import PreTrainingDataModule, MixtralConfig8x3B, MixtralModel
 
 from nemo.lightning.pytorch.optim.megatron import MegatronOptimizerModule as MegatronOptim, OptimizerConfig
@@ -148,4 +146,4 @@ def parse_args():
     return parser.parse_args()
 
 if __name__ == "__main__":
-    main(parse_args())
+    main(parse_args())