Apply isort and black reformatting

Signed-off-by: akoumpa <akoumpa@users.noreply.github.com>
NVIDIA · Aug 27, 2024 · 12b2655 · 12b2655
1 parent 3dd041f
commit 12b2655
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/nemo/collections/llm/gpt/model/mistral.py b/nemo/collections/llm/gpt/model/mistral.py
@@ -200,7 +200,8 @@ def convert_state(self, source, target):
             "output_layer.weight": "lm_head.weight",
         }
 
-        target.load_state_dict(StateDictTransformer(
+        target.load_state_dict(
+            StateDictTransformer(
                 source.state_dict(), mapping, transforms, transforms=[_export_qkv, _export_linear_fc1]
             )
         )
@@ -334,4 +335,4 @@ def _export_linear_fc1(linear_fc1):
     "MistralNeMo2407Config12B",
     "MistralNeMo2407Config123B",
     "MistralModel",
-]
+]
diff --git a/nemo/collections/llm/gpt/model/mixtral.py b/nemo/collections/llm/gpt/model/mixtral.py
@@ -372,11 +372,10 @@ def _export_moe_w1_w3(linear_fc1):
     return gate_proj, up_proj
 
 
-
 __all__ = [
     "MixtralConfig",
     "MixtralConfig8x3B",
     "MixtralConfig8x7B",
     "MixtralConfig8x22B",
     "MixtralModel",
-]
+]