Просмотр исходного кода

wasi_nn_tensorflowlite.cpp: reject non-fp32 input earlier (#4388)

this backend assumes fp32 here and there.
it's safer to reject unexpected inputs explicitly.
YAMAMOTO Takashi 8 месяцев назад
Родитель
Сommit
db7714f0f5
1 измененных файлов с 5 добавлено и 0 удалено
  1. 5 0
      core/iwasm/libraries/wasi-nn/src/wasi_nn_tensorflowlite.cpp

+ 5 - 0
core/iwasm/libraries/wasi-nn/src/wasi_nn_tensorflowlite.cpp

@@ -281,6 +281,11 @@ set_input(void *tflite_ctx, graph_execution_context ctx, uint32_t index,
 {
     TFLiteContext *tfl_ctx = (TFLiteContext *)tflite_ctx;
 
+    if (input_tensor->type != fp32) {
+        NN_ERR_PRINTF("unsupported input tensor type %u", input_tensor->type);
+        return runtime_error;
+    }
+
     wasi_nn_error res;
     if (success != (res = is_valid_graph_execution_context(tfl_ctx, ctx)))
         return res;