MLBEDSW-4635: yolo_v3 output diff
Fixed an issue where the scheduler would set the incorrect tensor
layout.
Signed-off-by: Jacob Bohlin <jacob.bohlin@arm.com>
Change-Id: I28abdf3f3c523d7da0cf8840316ece37dad364ab
diff --git a/ethosu/vela/scheduler.py b/ethosu/vela/scheduler.py
index 71007a3..dfb8867 100644
--- a/ethosu/vela/scheduler.py
+++ b/ethosu/vela/scheduler.py
@@ -351,7 +351,7 @@
if ps.primary_op:
# Set tensor format to NHCWB16 for output FeatureMaps, if possible
for output in ps.outputs:
- if output.purpose != TensorPurpose.FeatureMap:
+ if output in self.sg.output_tensors or output.purpose != TensorPurpose.FeatureMap:
continue
if not output.needs_linear_format:
output.set_format(TensorFormat.NHCWB16, arch)