From 8f142c6238155cf7d4f9a38a0048a81acc9c46f0 Mon Sep 17 00:00:00 2001
From: LukeAI <43993778+LukeAI@users.noreply.github.com>
Date: Wed, 23 Nov 2022 10:52:16 +0000
Subject: [PATCH] Update torch2onnx.py

Fixes https://github.com/xuanandsix/CLRNet-onnxruntime-and-tensorrt-demo/issues/3
---
 torch2onnx.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torch2onnx.py b/torch2onnx.py
index 377289d..58cfefc 100644
--- a/torch2onnx.py
+++ b/torch2onnx.py
@@ -25,7 +25,7 @@ def main():
         namekey = k[7:]    # 去掉module前缀
         new_state_dict[namekey] = v
         
-    net.load_state_dict(new_state_dict)
+    net.load_state_dict(new_state_dict, strict=False)
     
     dummy_input = torch.randn(1, 3, 320,800, device='cpu')
     torch.onnx.export(net, dummy_input, 'tusimple_r18.onnx',