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

Fix unguarded data type conversion from int64 to int in arg_min_max.cc

Browse Source
This commit is contained in:
Sidong-Wei 2020-11-16 17:26:58 -05:00
parent d267b0106b
commit 72cb2eff53
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorflow/lite/kernels/arg_min_max.cc
View File

@ -37,7 +37,12 @@ constexpr int kOutputTensor = 0;
TfLiteStatus ResizeOutput(TfLiteContext* context, const TfLiteTensor* input,
const TfLiteTensor* axis, TfLiteTensor* output) {
int axis_value = *GetTensorData<int>(axis);
int axis_value;
if (axis->type == kTfLiteInt64) {
axis_value = *GetTensorData<int64_t>(axis);
} else {
axis_value = *GetTensorData<int>(axis);
}
if (axis_value < 0) {
axis_value += NumDimensions(input);
}