Update add_nms.py

kweston · Jul 23, 2022 · aae7070 · aae7070
1 parent a5b853c
commit aae7070
Showing 1 changed file with 6 additions and 2 deletions.
diff --git 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
-import onnx_graphsurgeon as gs
+try:
+    import onnx_graphsurgeon as gs
+except Exception as e:
+    print('Import onnx_graphsurgeon failure: %s' % e)
+
 import logging
 
 LOGGER = logging.getLogger(__name__)
@@ -148,4 +152,4 @@ def save(self, output_path):
         self.graph.cleanup().toposort()
         model = gs.export_onnx(self.graph)
         onnx.save(model, output_path)
-        LOGGER.info(f"Saved ONNX model to {output_path}")
+        LOGGER.info(f"Saved ONNX model to {output_path}")