cleanup

Signed-off-by: Sangkug Lym <[email protected]>
NVIDIA · Nov 6, 2024 · 4bb1812 · 4bb1812
1 parent 2231cc8
commit 4bb1812
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py
@@ -286,7 +286,7 @@ def gemm(
             ), "SPLIT_PIPELINED_RS requires extra output tensor"
             # Disable the overlap between GEMM chunks at ampere and below
             major, _ = torch.cuda.get_device_capability()
-            overlap_gemm_chunks = True if major >= 9 else False
+            overlap_gemm_chunks = major >= 9
             args = tuple(
                 args
                 + (