Apply isort and black reformatting

Signed-off-by: akoumpa <[email protected]>
NVIDIA · Jul 26, 2024 · 12ce03c · 12ce03c
1 parent ce6b090
commit 12ce03c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/nemo/collections/llm/gpt/model/mixtral.py b/nemo/collections/llm/gpt/model/mixtral.py
@@ -123,7 +123,7 @@ def config(self) -> MixtralConfig8x7B:
             hidden_size=config.hidden_size,
             ffn_hidden_size=config.intermediate_size,
             max_position_embeddings=config.max_position_embeddings,  # TODO
-            seq_length=4096, #config.max_position_embeddings,
+            seq_length=4096,  # config.max_position_embeddings,
             # RoPE
             position_embedding_type='rope',
             rotary_base=config.rope_theta,
@@ -143,7 +143,7 @@ def config(self) -> MixtralConfig8x7B:
             # CPU init
             use_cpu_initialization=True,
             perform_initialization=False,
-            params_dtype=getattr(config, "torch_dtype", torch.bfloat16)
+            params_dtype=getattr(config, "torch_dtype", torch.bfloat16),
         )
 
 

diff --git a/nemo/lightning/io/state.py b/nemo/lightning/io/state.py
@@ -4,8 +4,8 @@
 from typing import Any, Callable, Dict, Generic, List, Optional, Tuple, TypeVar, Union, overload
 
 import numpy as np
-from torch import nn
 import torch
+from torch import nn
 
 SourceModuleT = TypeVar("SourceModuleT", bound=nn.Module)
 TargetModuleT = TypeVar("TargetModuleT", bound=nn.Module)
@@ -19,6 +19,7 @@ class TransformCTX:
     target: nn.Module
     target_state: dict
 
+
 @torch.no_grad
 def apply_transforms(
     source: nn.Module,