fix text_generator_main run error

google-ai-edge · Sep 24, 2024 · 39a3b07 · 39a3b07
1 parent c9973d2
commit 39a3b07
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/ai_edge_torch/generative/examples/cpp/text_generator_main.cc b/ai_edge_torch/generative/examples/cpp/text_generator_main.cc
@@ -154,6 +154,10 @@ tflite::SignatureRunner* GetSignatureRunner(
     std::map<std::string, std::vector<float>>& kv_cache) {
   tflite::SignatureRunner* runner =
       interpreter->GetSignatureRunner(signature_name.c_str());
+  int64_t flag = 0;
+  flag |= kTfLiteCustomAllocationFlagsNone;
+  flag |= kTfLiteCustomAllocationFlagsSkipAlignCheck;
+
   for (auto& [name, cache] : kv_cache) {
     TfLiteCustomAllocation allocation = {
         .data = static_cast<void*>(cache.data()),
@@ -162,9 +166,9 @@ tflite::SignatureRunner* GetSignatureRunner(
     // delegates support this in-place update. For those cases, we need to do
     // a ping-pong buffer and update the pointers between inference calls.
     TFLITE_MINIMAL_CHECK(runner->SetCustomAllocationForInputTensor(
-                             name.c_str(), allocation) == kTfLiteOk);
+                             name.c_str(), allocation, flag) == kTfLiteOk);
     TFLITE_MINIMAL_CHECK(runner->SetCustomAllocationForOutputTensor(
-                             name.c_str(), allocation) == kTfLiteOk);
+                             name.c_str(), allocation, flag) == kTfLiteOk);
   }
   TFLITE_MINIMAL_CHECK(runner->AllocateTensors() == kTfLiteOk);
   return runner;