Update add_nms.py

2022-07-23 04:54:18 +03:00 · 2022-07-23 04:54:18 +03:00 · aae70703f7
commit aae70703f7
parent a5b853c3fa
1 changed files with 6 additions and 2 deletions
--- a/utils/add_nms.py
+++ b/utils/add_nms.py
@ -1,7 +1,11 @@
 import numpy as np
 import onnx
 from onnx import shape_inference
+try:
    import onnx_graphsurgeon as gs
+except Exception as e:
+    print('Import onnx_graphsurgeon failure: %s' % e)
+
 import logging

 LOGGER = logging.getLogger(__name__)