[executorch][flat_tensor] update flat tensor header test

after changes in D68578075 Differential Revision: [D68971340](https://our.internmc.facebook.com/intern/diff/D68971340/) [ghstack-poisoned]
pytorch · Jan 31, 2025 · f9114f2 · f9114f2
1 parent 3413971
commit f9114f2
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/extension/flat_tensor/serialize/flat_tensor_header.cpp b/extension/flat_tensor/serialize/flat_tensor_header.cpp
@@ -76,7 +76,8 @@ uint64_t GetUInt64LE(const uint8_t* data) {
   if (size < FlatTensorHeader::kNumHeadBytes) {
     return Error::InvalidArgument;
   }
-  const uint8_t* header = reinterpret_cast<const uint8_t*>(data);
+  const uint8_t* header =
+      reinterpret_cast<const uint8_t*>(data) + kHeaderOffset;
 
   // Check magic bytes.
   if (std::memcmp(

diff --git a/extension/flat_tensor/serialize/flat_tensor_header.h b/extension/flat_tensor/serialize/flat_tensor_header.h
@@ -23,6 +23,12 @@ struct FlatTensorHeader {
    */
   static constexpr size_t kNumHeadBytes = 64;
 
+  /**
+   * The offset into the serialized FlatTensor data where the FlatTensor
+   * header should begin.
+   */
+  static constexpr size_t kHeaderOffset = 8;
+
   /**
    * The magic bytes that identify the header. This should be in sync with
    * the magic in executorch/extension/flat_tensor/serialize/serialize.py

diff --git a/extension/flat_tensor/test/flat_tensor_header_test.cpp b/extension/flat_tensor/test/flat_tensor_header_test.cpp
@@ -68,7 +68,7 @@ std::vector<uint8_t> CreateExampleFlatTensorHeader() {
   memset(ret.data(), 0x55, ret.size());
   // Copy the example header into the right offset.
   memcpy(
-      ret.data(),
+      ret.data() + FlatTensorHeader::kHeaderOffset,
       kExampleHeaderData,
       sizeof(kExampleHeaderData));
   return ret;