Fix fp16 ONNX detection for decoder models (#1276)

fix fp16 detection
huggingface · Aug 11, 2023 · a86f334 · a86f334
1 parent 3dfae0c
commit a86f334
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum/onnxruntime/base.py b/optimum/onnxruntime/base.py
@@ -152,7 +152,7 @@ def __init__(
 
         self.use_fp16 = False
         for inp in session.get_inputs():
-            if inp.name == "past_key_values" and inp.type == "tensor(float16)":
+            if "past_key_values" in inp.name and inp.type == "tensor(float16)":
                 self.use_fp16 = True
                 break