[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Signed-off-by: Jaime Cardenas <jaime@evolutionaryscale.ai>

Author: pre-commit-ci[bot]

tests/pytorch/selective_layernorm_mlp/test_cuda_graphs.py

@@ -168,6 +168,7 @@ def forward(self, input_: torch.Tensor, **kwargs) -> torch.Tensor:
 # Supported modules
 _test_cuda_graphs_modules: List[str] = ["selective_layernorm_mlp"]
 
+
 def _test_cuda_graphs(
     *,
     graph_mode: str,
@@ -200,7 +201,7 @@ def _test_cuda_graphs(
                 )
                 for _ in range(num_layers)
             ]
-        
+
         else:
             raise ValueError(f"Unknown module type ({module})")
 

tests/pytorch/selective_layernorm_mlp/test_deferred_init.py

@@ -7,7 +7,9 @@
 
 import transformer_engine.pytorch as te
 
-_core_modules = [te.SelectiveLayerNormMLP,]
+_core_modules = [
+    te.SelectiveLayerNormMLP,
+]
 _composed_modules = []
 
 batch_size = 32
@@ -16,6 +18,7 @@
 head_dim = 64
 dtype = torch.bfloat16
 
+
 class TestDeferredInit:
 
     @staticmethod