add version param

zhiqwang · Feb 23, 2022 · ef58e0c · ef58e0c
1 parent 3af7e4d
commit ef58e0c
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tools/export_model.py b/tools/export_model.py
@@ -116,9 +116,10 @@ def export_tensorrt(
     checkpoint_path: str,
     score_thresh: float,
     nms_thresh: float,
+    version: str,
     onnx_path: str,
     engine_path: str,
-    input_sample: torch.Tensor,
+    input_sample: torch.Tensor = None,
     detections_per_img: int = 100,
     workspace: int = 12,
 ):
@@ -128,6 +129,7 @@ def export_tensorrt(
         checkpoint_path,
         score_thresh=score_thresh,
         nms_thresh=nms_thresh,
+        version=version,
         onnx_path=onnx_path,
         engine_path=engine_path,
         input_sample=input_sample,
@@ -177,6 +179,7 @@ def cli_main():
             checkpoint_path,
             score_thresh=args.score_thresh,
             nms_thresh=args.nms_thresh,
+            version=args.version,
             onnx_path=str(onnx_path),
             engine_path=str(tensorrt_path),
             input_sample=input_sample,