Skip to content

Commit 3b63b4c

Browse files
committed
Apply isort and black reformatting
Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>
1 parent 6e4bd6f commit 3b63b4c

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

scripts/checkpoint_converters/convert_mamba2_pyt_to_nemo.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@
1616
import re
1717
from argparse import ArgumentParser
1818
from collections import defaultdict
19+
from datetime import timedelta
1920

21+
import megatron.core.parallel_state as ps
2022
import torch
2123
import torch.distributed as dist
2224
from megatron.core.dist_checkpointing.serialization import load_plain_tensors
@@ -25,10 +27,9 @@
2527
from megatron.core.transformer.spec_utils import import_module
2628
from megatron.training.arguments import core_transformer_config_from_args
2729
from omegaconf.omegaconf import OmegaConf
28-
from datetime import timedelta
29-
import megatron.core.parallel_state as ps
3030
from torch._C._distributed_c10d import PrefixStore
3131
from torch.distributed import rendezvous
32+
3233
from nemo.collections.nlp.models.language_modeling.megatron_mamba_model import MegatronMambaModel
3334
from nemo.collections.nlp.parts.megatron_trainer_builder import MegatronLMPPTrainerBuilder
3435
from nemo.collections.nlp.parts.utils_funcs import torch_dtype_from_precision
@@ -95,6 +96,7 @@ def get_args():
9596
args = parser.parse_args()
9697
return args
9798

99+
98100
try:
99101

100102
class Utils:

0 commit comments

Comments
 (0)