diff --git a/demo/__init__.py b/demo/__init__.py
new file mode 100644
index 00000000..f2699c4f
--- /dev/null
+++ b/demo/__init__.py
@@ -0,0 +1,3 @@
+from .predict import detect
+
+__all__ = ['detect']
diff --git a/mindyolo/utils/metrics.py b/mindyolo/utils/metrics.py
index 83c9e6ea..1565b7ab 100644
--- a/mindyolo/utils/metrics.py
+++ b/mindyolo/utils/metrics.py
@@ -70,7 +70,7 @@ def non_max_suppression(
         # Detections matrix nx6 (xyxy, conf, cls)
         if multi_label:
             i, j = (x[:, 5:] > conf_thres).nonzero()
-            x = np.concatenate((box[i], x[i, j + 5, None], j[:, None].astype(np.float)), 1)
+            x = np.concatenate((box[i], x[i, j + 5, None], j[:, None].astype(np.float32)), 1)
         else:  # best class only
             conf, j = x[:, 5:].max(1, keepdim=True)
             x = np.concatenate((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
diff --git a/tests/modules/test_deploy_predict.py b/tests/modules/test_demo_predict.py
similarity index 71%
rename from tests/modules/test_deploy_predict.py
rename to tests/modules/test_demo_predict.py
index e2da271f..676eb6c8 100644
--- a/tests/modules/test_deploy_predict.py
+++ b/tests/modules/test_demo_predict.py
@@ -2,27 +2,38 @@
 
 sys.path.append(".")
 import os
-import numpy as np
+import pytest
 from download import download
 import cv2
 
-from mindspore import nn
-
-from mindyolo.utils.metrics import non_max_suppression, scale_coords, xyxy2xywh
+from mindyolo.utils.config import load_config, Config
 from mindyolo.data import COCO80_TO_COCO91_CLASS
-from deploy.infer_engine.lite import LiteModel
 from mindyolo.utils.utils import draw_result
-from deploy.predict import detect
+from demo.predict import detect
+from mindyolo.models.model_factory import create_model
+
 
-def test_deploy_predict():
+@pytest.mark.parametrize("yaml_name", ['yolov5n.yaml'])
+def test_demo_predict(yaml_name):
     image_url = 'https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/image_cat.zip'
     path = download(image_url, './', kind="zip", replace=True)
     image_path = ('./image_cat/jpg/000000039769.jpg')
 
-    model_url = 'https://download.mindspore.cn/toolkits/mindyolo/yolov5/yolov5n_300e_mAP273-9b16bd7b-bd03027b.mindir'
-    model_path = download(model_url, './yolov5n.mindir', kind="file", replace=True)
-
-    network = LiteModel(model_path)
+    model_url = 'https://download.mindspore.cn/toolkits/mindyolo/yolov5/yolov5n_300e_mAP273-9b16bd7b.ckpt'
+    weight = download(model_url, './yolov5n.ckpt', kind="file", replace=True)
+    parent_dir = yaml_name[:6]
+    yaml_path = os.path.join('./configs', parent_dir, yaml_name)
+    cfg, _, _ = load_config(yaml_path)
+    cfg = Config(cfg)
+
+    network = create_model(
+        model_name=cfg.network.model_name,
+        model_cfg=cfg.network,
+        num_classes=cfg.data.nc,
+        sync_bn=False,
+        checkpoint_path=weight,
+    )
+    network.set_train(False)
     img = cv2.imread(image_path)
     result_dict = detect(
         network=network,
@@ -49,4 +60,4 @@ def test_deploy_predict():
 
 
 if __name__ == '__main__':
-    test_deploy_predict()
+    test_demo_predict('yolov5n.yaml')