diff --git a/wekws/bin/export_onnx.py b/wekws/bin/export_onnx.py
index 89bf002..75eb5b0 100644
--- a/wekws/bin/export_onnx.py
+++ b/wekws/bin/export_onnx.py
@@ -17,6 +17,7 @@ import argparse
 import torch
 import yaml
 
+import onnx
 import onnxruntime as ort
 
 from wekws.model.kws_model import init_model
diff --git a/wekws/model/kws_model.py b/wekws/model/kws_model.py
index 7e44bb4..92b76ab 100644
--- a/wekws/model/kws_model.py
+++ b/wekws/model/kws_model.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import sys
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
 import torch.nn as nn