onnx模式释放模型降低占用

douzero/evaluation/deep_agent.py

@@ -30,7 +30,8 @@ def _load_model(position, model_path, model_type, use_legacy, use_lite, use_onnx
     # torch.save(model.state_dict(), model_path.replace(".ckpt", "_nobn.ckpt"))
     model.eval()
     model_path = model_path + '.onnx'
-    if use_onnx and not os.path.exists(model_path):
+    if use_onnx:
+        if not os.path.exists(model_path):
             onnx_params = model.get_onnx_params(torch.device('cpu'))
             torch.onnx.export(
                 model,
@@ -43,6 +44,7 @@ def _load_model(position, model_path, model_type, use_legacy, use_lite, use_onnx
                 output_names=onnx_params['output_names'],
                 dynamic_axes=onnx_params['dynamic_axes']
             )
+        return None
 
     return model