Check and report allocation failures in InitializeRuntimeTensor.

PiperOrigin-RevId: 306497434
Change-Id: I1b94e40f5ea12e5f2c9fbb73be534b12016e76fe
This commit is contained in:
Robert David 2020-04-14 12:58:54 -07:00 committed by TensorFlower Gardener
parent 6f1bbf2cc2
commit e2cec54297

View File

@ -351,12 +351,29 @@ TfLiteStatus InitializeRuntimeTensor(
reinterpret_cast<TfLiteAffineQuantization*>( reinterpret_cast<TfLiteAffineQuantization*>(
allocator->AllocateFromTail(sizeof(TfLiteAffineQuantization), allocator->AllocateFromTail(sizeof(TfLiteAffineQuantization),
alignof(TfLiteAffineQuantization))); alignof(TfLiteAffineQuantization)));
if (quantization == nullptr) {
TF_LITE_REPORT_ERROR(error_reporter,
"Unable to allocate TfLiteAffineQuantization.\n");
return kTfLiteError;
}
quantization->zero_point = quantization->zero_point =
reinterpret_cast<TfLiteIntArray*>(allocator->AllocateFromTail( reinterpret_cast<TfLiteIntArray*>(allocator->AllocateFromTail(
TfLiteIntArrayGetSizeInBytes(channels), alignof(TfLiteIntArray))); TfLiteIntArrayGetSizeInBytes(channels), alignof(TfLiteIntArray)));
if (quantization->zero_point == nullptr) {
TF_LITE_REPORT_ERROR(error_reporter,
"Unable to allocate quantization->zero_point.\n");
return kTfLiteError;
}
quantization->scale = reinterpret_cast<TfLiteFloatArray*>( quantization->scale = reinterpret_cast<TfLiteFloatArray*>(
allocator->AllocateFromTail(TfLiteFloatArrayGetSizeInBytes(channels), allocator->AllocateFromTail(TfLiteFloatArrayGetSizeInBytes(channels),
alignof(TfLiteFloatArray))); alignof(TfLiteFloatArray)));
if (quantization->scale == nullptr) {
TF_LITE_REPORT_ERROR(error_reporter,
"Unable to allocate quantization->scale.\n");
return kTfLiteError;
}
quantization->zero_point->size = channels; quantization->zero_point->size = channels;
quantization->scale->size = channels; quantization->scale->size = channels;
int* zero_point_data = quantization->zero_point->data; int* zero_point_data = quantization->zero_point->data;