support attn_mask fp16 type

rui-ren · rui-ren · commit bacc05a69da5 · 2023-09-29T19:16:18.000Z
diff --git a/torch/onnx/symbolic_opset14.py b/torch/onnx/symbolic_opset14.py
@@ -190,7 +190,7 @@ def scaled_dot_product_attention(
         mul_qk_add = g.op("Add", mul_qk, attn_mask)
     elif (
         _type_utils.JitScalarType.from_value(attn_mask)
-        == _type_utils.JitScalarType.FLOAT
+        in (_type_utils.JitScalarType.FLOAT, _type_utils.JitScalarType.HALF)
     ):
         mul_qk_add = g.op("Add", mul_qk, attn_mask)
     else:
