From d0b629232bec498b422fe86ae95c0090a6aabfef Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Tue, 11 Jul 2023 15:57:06 -0400 Subject: [PATCH] Try slicing outputs differently --- src/deepsparse/utils/data.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/deepsparse/utils/data.py b/src/deepsparse/utils/data.py index b7fa1350dc..7200d55e56 100644 --- a/src/deepsparse/utils/data.py +++ b/src/deepsparse/utils/data.py @@ -249,6 +249,8 @@ def join_engine_outputs( # If we can't evenly divide with batch size, remove the remainder as padding if candidate_output[0].shape[0] > orig_batch_size: - candidate_output = [out[:orig_batch_size] for out in candidate_output] + for i in range(len(candidate_output)): + candidate_output[i] = candidate_output[i][:orig_batch_size] + return candidate_output