diff --git a/core/iwasm/libraries/wasi-nn/wasi_nn.c b/core/iwasm/libraries/wasi-nn/wasi_nn.c
index 7e9034708..48e2bcad0 100644
--- a/core/iwasm/libraries/wasi-nn/wasi_nn.c
+++ b/core/iwasm/libraries/wasi-nn/wasi_nn.c
@@ -7,7 +7,7 @@
 
 #include "wasm_export.h"
 
-#include "lib_run_inference.hpp"
+#include "wasi_nn_tensorflow.hpp"
 
 /**
  * @brief loader of tensorflow
diff --git a/core/iwasm/libraries/wasi-nn/lib_run_inference.cpp b/core/iwasm/libraries/wasi-nn/wasi_nn_tensorflow.cpp
similarity index 89%
rename from core/iwasm/libraries/wasi-nn/lib_run_inference.cpp
rename to core/iwasm/libraries/wasi-nn/wasi_nn_tensorflow.cpp
index 89901f56f..ee193baeb 100644
--- a/core/iwasm/libraries/wasi-nn/lib_run_inference.cpp
+++ b/core/iwasm/libraries/wasi-nn/wasi_nn_tensorflow.cpp
@@ -1,4 +1,4 @@
-#include "lib_run_inference.hpp"
+#include "wasi_nn_tensorflow.hpp"
 
 #include <tensorflow/lite/interpreter.h>
 #include <tensorflow/lite/kernels/register.h>
@@ -28,7 +28,7 @@ _load(graph_builder_array builder, graph_encoding encoding)
 
     if (model == nullptr) {
         printf("failure: null model \n");
-        return invalid_argument;
+        return missing_memory;
     }
 
     // Build the interpreter with the InterpreterBuilder.
@@ -38,7 +38,7 @@ _load(graph_builder_array builder, graph_encoding encoding)
 
     if (interpreter == nullptr) {
         printf("failure: null interpreter \n");
-        return invalid_argument;
+        return missing_memory;
     }
 
     return success;
@@ -49,8 +49,14 @@ _set_input(tensor input_tensor)
 {
     auto *input = interpreter->typed_input_tensor<float>(0);
 
+    if (input == nullptr) {
+        return missing_memory;
+    }
+
     for (int i = 0; i < input_tensor.dimensions[0]; i++) {
         input[i] = (float)input_tensor.data[i];
     }
+
     return success;
-}
+    
+}
\ No newline at end of file
diff --git a/core/iwasm/libraries/wasi-nn/lib_run_inference.hpp b/core/iwasm/libraries/wasi-nn/wasi_nn_tensorflow.hpp
similarity index 78%
rename from core/iwasm/libraries/wasi-nn/lib_run_inference.hpp
rename to core/iwasm/libraries/wasi-nn/wasi_nn_tensorflow.hpp
index 7ee1de205..b068aa312 100644
--- a/core/iwasm/libraries/wasi-nn/lib_run_inference.hpp
+++ b/core/iwasm/libraries/wasi-nn/wasi_nn_tensorflow.hpp
@@ -1,5 +1,5 @@
-#ifndef LIB_RUN_INFERENCE_HPP
-#define LIB_RUN_INFERENCE_HPP
+#ifndef WASI_NN_TENSORFLOW_HPP
+#define WASI_NN_TENSORFLOW_HPP
 
 #include <stdio.h>