From 706e9c1fd097846f20cac1900cee239f68e2b159 Mon Sep 17 00:00:00 2001 From: Alexandros Koumparoulis Date: Fri, 27 Sep 2024 09:54:27 -0700 Subject: [PATCH] bot happy Signed-off-by: Alexandros Koumparoulis --- tests/collections/llm/megatron_mixtral_pretraining.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/collections/llm/megatron_mixtral_pretraining.py b/tests/collections/llm/megatron_mixtral_pretraining.py index 228015abf324..f973f51bac18 100644 --- 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()) \ No newline at end of file + main(parse_args())