From 2f496f673ef3c835e3076c328e9eac122b510273 Mon Sep 17 00:00:00 2001 From: inisis Date: Thu, 16 May 2024 11:56:10 +0000 Subject: [PATCH] fix bugs when constant_values are list Signed-off-by: inisis --- .../onnx-graphsurgeon/onnx_graphsurgeon/ir/graph.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tools/onnx-graphsurgeon/onnx_graphsurgeon/ir/graph.py b/tools/onnx-graphsurgeon/onnx_graphsurgeon/ir/graph.py index b40968f4..6ffbb594 100644 --- a/tools/onnx-graphsurgeon/onnx_graphsurgeon/ir/graph.py +++ b/tools/onnx-graphsurgeon/onnx_graphsurgeon/ir/graph.py @@ -1282,15 +1282,20 @@ def should_eval_foldable(tensor): # No need to fold tensors that are already constant. continue - if size_threshold is not None and values.nbytes > size_threshold: + if isinstance(values, list): + nbytes = sum([v.nbytes for v in values]) + else: + nbytes = values.nbytes + + if size_threshold is not None and nbytes > size_threshold: G_LOGGER.debug( "Will not fold: '{:}' since its size in bytes ({:}) exceeds the size threshold ({:})".format( - name, values.nbytes, size_threshold + name, nbytes, size_threshold ) ) continue - elif size_threshold is None and values.nbytes > (1 << 20): - large_tensors[name] = values.nbytes + elif size_threshold is None and nbytes > (1 << 20): + large_tensors[name] = nbytes tensor.to_constant(values) tensor.inputs.clear() # Constants do not need inputs