diff --git a/tensorflow/lite/delegates/gpu/common/object_reader.cc b/tensorflow/lite/delegates/gpu/common/object_reader.cc index 55a0aea01a1..f299232a8a7 100644 --- a/tensorflow/lite/delegates/gpu/common/object_reader.cc +++ b/tensorflow/lite/delegates/gpu/common/object_reader.cc @@ -37,7 +37,7 @@ absl::Status ObjectReader::ReadNonConstantTensor( } if (tensor_to_value->find(tensor_idx) == tensor_to_value->end()) { - const TfLiteTensor& tflite_tensor = context->tensors[tensor_idx]; + TfLiteTensor& tflite_tensor = context->tensors[tensor_idx]; if (tflite::IsConstantTensor(&tflite_tensor)) { return absl::InvalidArgumentError(absl::StrCat( "ReadNonConstantTensor: value is a constant tensor: ", tensor_idx)); @@ -58,6 +58,7 @@ absl::Status ObjectReader::ReadNonConstantTensor( &fp_tensor_index) != kTfLiteOk) { return absl::InternalError("Could not add new tensor to graph"); } + // Remember this tensor for later. (*quant_conversion_map)[fp_tensor_index] = tensor_idx; (*quant_conversion_map)[tensor_idx] = fp_tensor_index; @@ -67,6 +68,9 @@ absl::Status ObjectReader::ReadNonConstantTensor( ConvertTfLiteTensorToTensorRef(*fp_tflite_tensor, &value->tensor)); value->tensor.ref = fp_tensor_index; value->quant_params.emplace(); + // tflite_tensor from the outer scope is invalidated due to calling + // CreateNewTensorWithDifferentType + tflite_tensor = context->tensors[tensor_idx]; RETURN_IF_ERROR( PopulateQuantParams(tflite_tensor, &value->quant_params.value())); (*tensor_to_value)[fp_tensor_index] = value; diff --git a/tensorflow/lite/delegates/utils.cc b/tensorflow/lite/delegates/utils.cc index 873cadc180f..289586c5346 100644 --- a/tensorflow/lite/delegates/utils.cc +++ b/tensorflow/lite/delegates/utils.cc @@ -29,8 +29,8 @@ TfLiteStatus CreateNewTensorWithDifferentType(TfLiteContext* context, TfLiteType new_type, TfLiteTensor** new_tensor, int* new_tensor_index) { - const TfLiteTensor& original_tensor = context->tensors[original_tensor_index]; TF_LITE_ENSURE_STATUS(context->AddTensors(context, 1, new_tensor_index)); + const TfLiteTensor& original_tensor = context->tensors[original_tensor_index]; *new_tensor = &context->tensors[*new_tensor_index]; (*new_tensor)->type = new_type; (*new_tensor)->allocation_type = kTfLiteArenaRw; diff --git a/tensorflow/lite/delegates/utils.h b/tensorflow/lite/delegates/utils.h index 12684fcb84a..a9fb67316fc 100644 --- a/tensorflow/lite/delegates/utils.h +++ b/tensorflow/lite/delegates/utils.h @@ -33,7 +33,8 @@ namespace tflite { namespace delegates { // Creates a new Read/Write tensor having the same shape as the original, but -// with a different type. +// with a different type. Note that this might void existing references to +// tensors. TfLiteStatus CreateNewTensorWithDifferentType(TfLiteContext* context, const int original_tensor_index, TfLiteType new_type,