Skip to content

Commit

Permalink
fix isort
Browse files Browse the repository at this point in the history
Signed-off-by: jiemingz <jiemingz@nvidia.com>
  • Loading branch information
jiemingz committed Mar 22, 2024
1 parent bacb335 commit 482053f
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion nemo/collections/nlp/parts/nlp_overrides.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@

try:
from apex.transformer.pipeline_parallel.utils import get_num_microbatches

from nemo.core.optim.distributed_adam import MegatronDistributedFusedAdam

HAVE_APEX = True
Expand All @@ -87,12 +88,12 @@
try:
from megatron.core import dist_checkpointing, parallel_state
from megatron.core.dist_checkpointing.dict_utils import dict_list_map_outplace
from megatron.core.dist_checkpointing.strategies import tensorstore
from megatron.core.dist_checkpointing.optimizer import (
get_param_id_to_sharded_param_map,
make_sharded_optimizer_tensor,
optim_state_to_sharding_state,
)
from megatron.core.dist_checkpointing.strategies import tensorstore
from megatron.core.transformer.module import Float16Module as MCoreFloat16Module
from megatron.core.transformer.transformer_layer import TransformerLayer as MCoreTransformerLayer

Expand Down

0 comments on commit 482053f

Please sign in to comment.