fix lint

Signed-off-by: Alexandros Koumparoulis <[email protected]>

Author: akoumpa

nemo_automodel/_transformers/auto_model.py

@@ -23,14 +23,6 @@
 
 import torch
 from torch.nn.attention import SDPBackend, sdpa_kernel
-
-import nemo_automodel.components.distributed.utils as dist_utils
-from nemo_automodel import __version__
-from nemo_automodel._transformers.registry import ModelRegistry
-from nemo_automodel.components.distributed.init_utils import get_local_world_size_preinit, get_world_size_safe
-from nemo_automodel.components.utils.model_utils import resolve_trust_remote_code
-from nemo_automodel.shared.import_utils import safe_import
-from nemo_automodel.shared.utils import dtype_from_str
 from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
@@ -42,6 +34,14 @@
 from transformers.modeling_utils import _get_resolved_checkpoint_files
 from transformers.models.auto.auto_factory import _BaseAutoModelClass
 
+import nemo_automodel.components.distributed.utils as dist_utils
+from nemo_automodel import __version__
+from nemo_automodel._transformers.registry import ModelRegistry
+from nemo_automodel.components.distributed.init_utils import get_local_world_size_preinit, get_world_size_safe
+from nemo_automodel.components.utils.model_utils import resolve_trust_remote_code
+from nemo_automodel.shared.import_utils import safe_import
+from nemo_automodel.shared.utils import dtype_from_str
+
 HAS_LIGER_KERNEL, liger_kernel_trf = safe_import("liger_kernel.transformers")
 
 logger = logging.getLogger(__name__)

nemo_automodel/components/models/llama/model.py

@@ -33,10 +33,6 @@
 
 import torch
 import torch.nn as nn
-
-from nemo_automodel.components.models.common.combined_projection import CombinedGateUpMLP, CombinedQKVAttentionMixin
-from nemo_automodel.components.models.llama.state_dict_adapter import LlamaStateDictAdapter
-from nemo_automodel.shared.import_utils import get_check_model_inputs_decorator
 from transformers import LlamaConfig
 from transformers.cache_utils import Cache, DynamicCache
 from transformers.masking_utils import create_causal_mask
@@ -54,6 +50,10 @@
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, can_return_tuple
 
+from nemo_automodel.components.models.common.combined_projection import CombinedGateUpMLP, CombinedQKVAttentionMixin
+from nemo_automodel.components.models.llama.state_dict_adapter import LlamaStateDictAdapter
+from nemo_automodel.shared.import_utils import get_check_model_inputs_decorator
+
 check_model_inputs = get_check_model_inputs_decorator()