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 +}