From ef4d53818d720f4c1d742125fbb48eafe481fd21 Mon Sep 17 00:00:00 2001
From: WangChaofeng <syusuke0511@163.com>
Date: Thu, 27 May 2021 20:10:14 +0800
Subject: [PATCH]  ONNX export in .train() mode fix (#3362)

---
 models/export.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/models/export.py b/models/export.py
index 65721f65d888..0d1147938e37 100644
--- a/models/export.py
+++ b/models/export.py
@@ -97,6 +97,8 @@
             print(f'{prefix} starting export with onnx {onnx.__version__}...')
             f = opt.weights.replace('.pt', '.onnx')  # filename
             torch.onnx.export(model, img, f, verbose=False, opset_version=opt.opset_version, input_names=['images'],
+                              training=torch.onnx.TrainingMode.TRAINING if opt.train else torch.onnx.TrainingMode.EVAL,
+                              do_constant_folding=not opt.train,
                               dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # size(1,3,640,640)
                                             'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)