Update tools/demo.py

tools/demo.py

@@ -126,28 +126,15 @@ def export_model(runner,
     deploy_model(fake_input)
 
     os.makedirs('work_dirs', exist_ok=True)
-    save_onnx_path = os.path.join(
-        'work_dirs', 'yolow-l.onnx')
-    # export onnx
-    with BytesIO() as f:
-        output_names = ['num_dets', 'boxes', 'scores', 'labels']
-        torch.onnx.export(
-            deploy_model,
-            fake_input,
-            f,
-            input_names=['images'],
-            output_names=output_names,
-            opset_version=12)
-        f.seek(0)
-        onnx_model = onnx.load(f)
-        onnx.checker.check_model(onnx_model)
-    onnx_model, check = onnxsim.simplify(onnx_model)
-    onnx.save(onnx_model, save_onnx_path)
+    save_pt_path = os.path.join(
+        'work_dirs', 'yolow-l.pt')
+    
+    # Save model as a .pt file
+    torch.save(deploy_model.state_dict(), save_pt_path)
 
     del base_model
     del deploy_model
-    del onnx_model
-    return gr.update(visible=True), save_onnx_path
+    return gr.update(visible=True), save_pt_path
 
 
 def demo(runner, args, cfg):