Skip to content

Commit

Permalink
Fix TensorDataBuilder returning wrong offset
Browse files Browse the repository at this point in the history
When adding a new tensor, the returned offset was the end of the previous tensor
instead of the start of the new tensor. These offsets may differ if padding
bytes are added for alignment. The problem went unnoticed before because all the
supported data types had the same size (32bit), and that has now changed.
  • Loading branch information
robertknight committed Sep 6, 2024
1 parent d7c2040 commit 387cdb8
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion src/model_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -990,12 +990,14 @@ impl TensorDataBuilder {
let padding = offset.next_multiple_of(align) - offset;
self.data.extend(std::iter::repeat(0).take(padding));

let start_offset = self.data.len();

for x in data {
let bytes = x.to_le_bytes();
self.data.extend(bytes.as_ref());
}

offset
start_offset
}

/// Consume the builder and return the finalized tensor data buffer.
Expand Down

0 comments on commit 387cdb8

Please sign in to comment.