experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

TFLM: nit catch memory allocation failures

Browse Source 
PiperOrigin-RevId: 283699825
Change-Id: I0434f161c1ac6d5578e7a35f26f7a8344cc3cf6f
This commit is contained in:
Tiezhen WANG 2019-12-04 00:05:37 -08:00 committed by TensorFlower Gardener
parent 92cb2db3dc
commit 90a754c965
1 changed files with 11 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
tensorflow/lite/experimental/micro/micro_allocator.cc
View File

@ -89,6 +89,11 @@ MicroAllocator::MicroAllocator(TfLiteContext* context, const Model* model,
reinterpret_cast<TfLiteTensor*>(memory_allocator_.AllocateFromTail(
sizeof(TfLiteTensor) * context_->tensors_size,
alignof(TfLiteTensor)));
if (context_->tensors == nullptr) {
error_reporter_->Report(
"Failed to allocate memory for context->tensors, %d bytes required",
sizeof(TfLiteTensor) * context_->tensors_size);
}
// Null all inputs so we can later perform a null check to avoid re-allocating
// registered pre-allocated inputs.
@ -230,6 +235,12 @@ TfLiteStatus MicroAllocator::FinishTensorAllocation() {
TensorInfo* tensor_info =
reinterpret_cast<TensorInfo*>(tmp_allocator.AllocateFromTail(
sizeof(TensorInfo) * tensors_size, alignof(TensorInfo)));
if (tensor_info == nullptr) {
error_reporter_->Report(
"Failed to allocate memory for tensor_info, %d bytes required",
sizeof(TfLiteTensor) * context_->tensors_size);
return kTfLiteError;
}
// Set up the runtime data structures for all tensors.
for (size_t i = 0; i < tensors_size; ++i) {