Check and report allocation failures in InitializeRuntimeTensor.
PiperOrigin-RevId: 306497434 Change-Id: I1b94e40f5ea12e5f2c9fbb73be534b12016e76fe
This commit is contained in:
parent
6f1bbf2cc2
commit
e2cec54297
@ -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;
|
||||||
|
Loading…
Reference in New Issue
Block a user