diff options
Diffstat (limited to 'ethosu/vela/tflite_writer.py')
-rw-r--r-- | ethosu/vela/tflite_writer.py | 10 |
1 files changed, 6 insertions, 4 deletions
diff --git a/ethosu/vela/tflite_writer.py b/ethosu/vela/tflite_writer.py index 3f3b7b1b..92b5c6b0 100644 --- a/ethosu/vela/tflite_writer.py +++ b/ethosu/vela/tflite_writer.py @@ -414,12 +414,14 @@ class TFLiteSerialiser: if tens.mem_type in set((MemType.Scratch, MemType.Scratch_fast)) and tens.address is not None: offsets[idx] = np.int32(tens.address) - metadata_buffer = np.array([version, subgraph_idx, nbr_tensors] + offsets) - self.buffers_to_write.append(metadata_buffer) + self.nng.metadata.append(("OfflineMemoryAllocation", np.array([version, subgraph_idx, nbr_tensors] + offsets))) - buffers_offset = self.write_offset_vector([self.serialise_buffer(buf) for buf in self.buffers_to_write]) + metadata_list = [] + for name, buffer in self.nng.metadata: + self.buffers_to_write.append(buffer) + metadata_list.append((name, len(self.buffers_to_write) - 1)) - metadata_list = [("OfflineMemoryAllocation", len(self.buffers_to_write) - 1)] + buffers_offset = self.write_offset_vector([self.serialise_buffer(buf) for buf in self.buffers_to_write]) metadata_offset = self.write_offset_vector([self.serialise_metadata(metadata) for metadata in metadata_list]) Model.ModelStart(builder) |