From 750c42e43eb38bf23659fcee50576156acd86c77 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 19 Jan 2022 12:24:40 -1000 Subject: [PATCH] `export.py` automatic `forward_export` (#6352) * `export.py` automatic `forward_export` * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- export.py | 3 ++- tutorial.ipynb | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/export.py b/export.py index a7a79b46b8bb..589b381e035a 100644 --- a/export.py +++ b/export.py @@ -427,7 +427,8 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' elif isinstance(m, Detect): m.inplace = inplace m.onnx_dynamic = dynamic - # m.forward = m.forward_export # assign forward (optional) + if hasattr(m, 'forward_export'): + m.forward = m.forward_export # assign custom forward (optional) for _ in range(2): y = model(im) # dry runs diff --git a/tutorial.ipynb b/tutorial.ipynb index b160e75adb58..6ff20dc36c40 100644 --- a/tutorial.ipynb +++ b/tutorial.ipynb @@ -1099,4 +1099,4 @@ "outputs": [] } ] -} \ No newline at end of file +}