[MLBEDSW-3891] Fix reading back in an ethos-u custom op
Fixed assertion when reading back in an ethos-u custom op.
Signed-off-by: Fredrik Svedberg <fredrik.svedberg@arm.com>
Change-Id: I275ec9187ffead1e96f2522ecbd658328fa4ef69
diff --git a/ethosu/vela/tflite_writer.py b/ethosu/vela/tflite_writer.py
index 18905e3..82a063f 100644
--- a/ethosu/vela/tflite_writer.py
+++ b/ethosu/vela/tflite_writer.py
@@ -411,16 +411,19 @@
subgraph_idx = np.int32(len(self.subgraphs_to_write)) # Only 1 supported currently
nbr_tensors = np.int32(len(self.tensor_map))
- # An offset of -1 indicates that the tensor will be allocated online by Tensorflow Lite Micro
- offsets = [np.int32(-1)] * nbr_tensors
+ if not any([name == b"OfflineMemoryAllocation" for name, _ in self.nng.metadata]):
+ # An offset of -1 indicates that the tensor will be allocated online by Tensorflow Lite Micro
+ offsets = [np.int32(-1)] * nbr_tensors
- # Ensure that the order of the offsets match the order of the tensors
- for tens, idx in self.tensor_map.items():
- # Set offsets for tensor allocated in Tensor Arena or in the scratch_fast area
- if tens.mem_type in (MemType.Scratch, MemType.Scratch_fast):
- offsets[idx] = np.int32(tens.address) if tens.address is not None else np.int32(0)
+ # Ensure that the order of the offsets match the order of the tensors
+ for tens, idx in self.tensor_map.items():
+ # Set offsets for tensor allocated in Tensor Arena or in the scratch_fast area
+ if tens.mem_type in (MemType.Scratch, MemType.Scratch_fast):
+ offsets[idx] = np.int32(tens.address) if tens.address is not None else np.int32(0)
- self.nng.metadata.append(("OfflineMemoryAllocation", np.array([version, subgraph_idx, nbr_tensors] + offsets)))
+ self.nng.metadata.append(
+ ("OfflineMemoryAllocation", np.array([version, subgraph_idx, nbr_tensors] + offsets))
+ )
metadata_list = []
for name, buffer in self.nng.metadata: