diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 37122f0aa0858281676b43d8c525a113677b2375..6bcab270888965e2548bb9a733c851ee47b09a1f 100644 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -316,6 +316,10 @@ def get_output_tensors(args, mode, predictor): output_name = 'softmax_0.tmp_0' if output_name in output_names: return [predictor.get_output_handle(output_name)] + else: + for output_name in output_names: + output_tensor = predictor.get_output_handle(output_name) + output_tensors.append(output_tensor) else: for output_name in output_names: output_tensor = predictor.get_output_handle(output_name)