bug fix for using return_layernorm_output=True

the current implementation would release the output of ln, leading to an error if setting `return_layernorm_output=True`. 

Signed-off-by: Liyuan Liu <[email protected]>

transformer_engine/pytorch/module/layernorm_mlp.py

@@ -373,7 +373,7 @@ def forward(
                 ub=ub_obj_lnout if ub_overlap_ag else None,
                 extra_output_tensor=ln_out if ub_overlap_ag else None,
             )
-            if not is_grad_enabled:
+            if not is_grad_enabled and not return_layernorm_output:
                 clear_tensor_data(ln_out_total)
 
             if bias_gelu_nvfusion: