Skip to content

Commit b15beb9

Browse files
authored
unchecked pointer fix (#3058)
@tensorflow/micro Handle the failure to allocate temporary memory for a TfLiteTensor in a more graceful way. This allows most kernels in the Prepare phase to not crash due to a DCHECK on memory allocation failure. bug=fixes #3057
1 parent 39209b4 commit b15beb9

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

tensorflow/lite/micro/micro_allocator.cc

+6-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
1+
/* Copyright 2025 The TensorFlow Authors. All Rights Reserved.
22
33
Licensed under the Apache License, Version 2.0 (the "License");
44
you may not use this file except in compliance with the License.
@@ -871,6 +871,11 @@ TfLiteTensor* MicroAllocator::AllocateTempTfLiteTensor(
871871
TfLiteTensor* tensor = reinterpret_cast<TfLiteTensor*>(
872872
non_persistent_buffer_allocator_->AllocateTemp(sizeof(TfLiteTensor),
873873
alignof(TfLiteTensor)));
874+
if (tensor == nullptr) {
875+
MicroPrintf("Failed to allocate temp. memory for tensor %d, subgraph %d",
876+
tensor_index, subgraph_index);
877+
return nullptr;
878+
}
874879

875880
// Populate any fields from the flatbuffer, since this TfLiteTensor struct is
876881
// allocated in the temp section of the arena, ensure that additional

0 commit comments

Comments
 (0)