Fix incorrect dtype in LayerNormLinear (#483)

Signed-off-by: Tim Moon <[email protected]> Co-authored-by: Kirthi Shankar Sivamani <[email protected]>
NVIDIA · Oct 20, 2023 · 1afb625 · 1afb625
1 parent d097883
commit 1afb625
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -112,7 +112,7 @@ def forward(
             ub_obj_lnout = get_ub("qkv_fprop")
             ln_out = ub_obj_lnout.get_ubuf_output(0)
         else:
-            ln_out_dtype = torch.uint8 if fp8 else inputmat.dtype
+            ln_out_dtype = torch.uint8 if (fp8 and not return_layernorm_output) else inputmat.dtype
             ln_out = torch.empty_like(inputmat, dtype=ln_out_dtype)
         if ub_atomic_gemm_ag:
             assert fp8, "AtomicGemm overlap supported only for FP8 GEMM."