diff --git a/core/iwasm/libraries/wasi-nn/test/models/mult_outputs.py b/core/iwasm/libraries/wasi-nn/test/models/mult_outputs.py
index 98a50129c..8506e1806 100755
--- a/core/iwasm/libraries/wasi-nn/test/models/mult_outputs.py
+++ b/core/iwasm/libraries/wasi-nn/test/models/mult_outputs.py
@@ -3,7 +3,7 @@
 
 import tensorflow as tf
 import numpy as np
-from keras.layers import AveragePooling2D, Conv2D
+from tensorflow.keras.layers import AveragePooling2D, Conv2D
 
 from tensorflow.keras import Input, Model