diff options
| -rw-r--r-- | src/inference_engine_tflite.cpp | 7 |
1 files changed, 7 insertions, 0 deletions
diff --git a/src/inference_engine_tflite.cpp b/src/inference_engine_tflite.cpp index 8490b4d..d1a7687 100644 --- a/src/inference_engine_tflite.cpp +++ b/src/inference_engine_tflite.cpp @@ -225,6 +225,10 @@ namespace TFLiteImpl pBuff = (void *) mInterpreter->typed_tensor<uint8_t>( mOutputLayerId[idx]); buffer = { pBuff, INFERENCE_TENSOR_DATA_TYPE_UINT8, size, 1 }; + } else if (mInterpreter->tensor(mOutputLayerId[idx])->type == kTfLiteInt64) { + LOGI("type is kTfLiteInt64"); + pBuff = (void*)mInterpreter->typed_tensor<int64_t>(mOutputLayerId[idx]); + buffer = {pBuff, INFERENCE_TENSOR_DATA_TYPE_INT64, size * 8, 1}; } else if (mInterpreter->tensor(mOutputLayerId[idx])->type == kTfLiteFloat32) { LOGI("type is kTfLiteFloat32"); @@ -295,6 +299,9 @@ namespace TFLiteImpl if (mInterpreter->tensor((*iter))->type == kTfLiteUInt8) { LOGI("type is kTfLiteUInt8"); tensor_info.data_type = INFERENCE_TENSOR_DATA_TYPE_UINT8; + } else if (mInterpreter->tensor((*iter))->type == kTfLiteInt64) { + LOGI("type is kTfLiteInt64"); + tensor_info.data_type = INFERENCE_TENSOR_DATA_TYPE_INT64; } else if (mInterpreter->tensor((*iter))->type == kTfLiteFloat32) { LOGI("type is kTfLiteFloat32"); tensor_info.data_type = INFERENCE_TENSOR_DATA_TYPE_FLOAT32; |