Cleaned up the files per pylint

Signed-off-by: Oleg Goncharov <[email protected]>
NVIDIA · Feb 26, 2024 · 4f8e7a9 · 4f8e7a9
1 parent 68d8f78
commit 4f8e7a9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
@@ -2678,11 +2678,11 @@ def forward(
                 sequence_start:sequence_end, batch_start:batch_end, ...] = value_layer
             key_layer = inference_key_memory[:sequence_end, batch_start:batch_end, ...]
             value_layer = inference_value_memory[:sequence_end, batch_start:batch_end, ...]
-            
+
             if qkv_format == "bshd":
                 key_layer = key_layer.transpose(0, 1)
                 value_layer = value_layer.transpose(0, 1)
-            
+
             key_layer = key_layer.contiguous()
             value_layer = value_layer.contiguous()
 

diff --git a/transformer_engine/pytorch/softmax.py b/transformer_engine/pytorch/softmax.py
@@ -29,7 +29,7 @@ def _get_default_causal_mask(sq: int, sk: int) -> torch.Tensor:
     if matrix_shape not in _default_causal_mask:
         diagonal_offset = sk - sq + 1
         _default_causal_mask[matrix_shape] = torch.triu(
-            torch.ones(sq, sk, dtype=torch.bool, device="cuda"), 
+            torch.ones(sq, sk, dtype=torch.bool, device="cuda"),
             diagonal=diagonal_offset)
     return _default_causal_mask[matrix_shape]