force evaluated shape to int

huggingface · Jan 15, 2025 · aad9aaf · aad9aaf
1 parent c9b45ee
commit aad9aaf
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/optimum/onnxruntime/modeling_ort.py b/optimum/onnxruntime/modeling_ort.py
@@ -780,7 +780,7 @@ def _output_shape_inference(self, axis_name: Union[str, int], dimensions: Dict[s
             if token in dimensions:
                 tokens[idx] = str(dimensions[token])
 
-        return eval(" ".join(tokens))
+        return int(eval(" ".join(tokens)))
 
     # TODO: this method is bloated with state arguments (that are accesible using self) why ?
     def _prepare_io_binding(
@@ -859,6 +859,7 @@ def _prepare_io_binding(
                 output_shape = []
                 for axis_name in output_node.shape:
                     output_shape.append(self._output_shape_inference(axis_name, dimensions))
+
             output_buffer = self._prepare_output_buffer(model, output_shape, output_name)
 
             io_binding.bind_output(