瀏覽代碼

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 11 月之前
父節點
當前提交
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;