Try slicing outputs differently

neuralmagic · Jul 11, 2023 · d0b6292 · d0b6292
1 parent 64b45ac
commit d0b6292
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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