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

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

Author: pre-commit-ci[bot]

transformer_engine/pytorch/module/layernorm_linear.py

@@ -895,22 +895,13 @@ def wgrad_gemm(
                     del grad_bias_
 
                     # Deallocate input tensor if permitted
-                    if (
-                        not ctx.return_layernorm_output
-                        and not ctx.return_layernorm_output_gathered
-                    ):
+                    if not ctx.return_layernorm_output and not ctx.return_layernorm_output_gathered:
                         # Do not need to return layernorm output
                         clear_tensor_data(ln_out)
-                    elif (
-                        ctx.return_layernorm_output_gathered
-                        and ctx.ln_out_needs_gather
-                    ):
+                    elif ctx.return_layernorm_output_gathered and ctx.ln_out_needs_gather:
                         # ln_out is not the returned tensor
                         clear_tensor_data(ln_out)
-                    if (
-                        ctx.ln_out_needs_gather
-                        and not ctx.ub_bulk_dgrad
-                    ):
+                    if ctx.ln_out_needs_gather and not ctx.ub_bulk_dgrad:
                         clear_tensor_data(ln_out_total)
 
                 # Update grad input if overlapping reduce-scatter with wgrad GEMM

transformer_engine/pytorch/module/linear.py

@@ -890,7 +890,8 @@ def wgrad_gemm(
                         clear_tensor_data(inputmat_total)
 
                     if (
-                        ctx.parallel_mode == "row" and ctx.sequence_parallel
+                        ctx.parallel_mode == "row"
+                        and ctx.sequence_parallel
                         and not ctx.ub_overlap_ag
                     ):
                         clear_tensor_data(grad_output)