Fix bf16 te.Linear with no tokens.

NVIDIA · Feb 22, 2024 · 9163446 · 9163446
1 parent 142a349
commit 9163446
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py
@@ -176,6 +176,8 @@ def gemm(
         grad_bias = empty_tensor
 
     bias = bias if use_bias else empty_tensor
+    if A.nelement() == 0 or B.nelement() == 0:
+        return out, gelu_input
 
     assert A.dtype == dtype and B.dtype == dtype, \
         f'Expected dtype={dtype}, but found A.dtype={A.dtype} and B.dtype={B.dtype}'