diff options
Diffstat (limited to 'ethosu/vela/tflite_writer.py')
-rw-r--r-- | ethosu/vela/tflite_writer.py | 14 |
1 files changed, 8 insertions, 6 deletions
diff --git a/ethosu/vela/tflite_writer.py b/ethosu/vela/tflite_writer.py index d134c07c..7aab01f2 100644 --- a/ethosu/vela/tflite_writer.py +++ b/ethosu/vela/tflite_writer.py @@ -116,35 +116,35 @@ class TFLiteSerialiser: builder.StartVector(1, len(v), alignment) for e in v[::-1]: builder.PrependByte(e) - return builder.EndVector(len(v)) + return builder.EndVector() def write_int_vector(self, v): builder = self.builder builder.StartVector(4, len(v), 4) for e in v[::-1]: builder.PrependInt32(e) - return builder.EndVector(len(v)) + return builder.EndVector() def write_long_vector(self, v): builder = self.builder builder.StartVector(8, len(v), 8) for e in v[::-1]: builder.PrependInt64(e) - return builder.EndVector(len(v)) + return builder.EndVector() def write_float_vector(self, v): builder = self.builder builder.StartVector(4, len(v), 4) for e in v[::-1]: builder.PrependFloat32(e) - return builder.EndVector(len(v)) + return builder.EndVector() def write_offset_vector(self, v): builder = self.builder builder.StartVector(4, len(v), 4) for e in v[::-1]: builder.PrependUOffsetTRelative(e) - return builder.EndVector(len(v)) + return builder.EndVector() def assign_buffers_to_tensors(self, tensors, scratch_tensor): if scratch_tensor is not None: @@ -389,13 +389,15 @@ class TFLiteSerialiser: def write_aligned_bytes(self, buf): builder = self.builder + builder.assertNotNested() builder.nested = True data = bytes(buf) length_bytes = UOffsetTFlags.py_type(len(data)) + builder.vectorNumElems = length_bytes builder.Prep(16, length_bytes) # Reserve aligned storage builder.head = UOffsetTFlags.py_type(builder.Head() - length_bytes) # Update FlatBuffer internal pointer builder.Bytes[builder.Head() : builder.Head() + length_bytes] = data # Assign bytes to aligned area - return builder.EndVector(length_bytes) + return builder.EndVector() def serialise_buffer(self, buf): builder = self.builder |