From 4fb8cb353f7d1945e3e1b270980c883c82297d2f Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 24 Aug 2020 21:59:26 -0700 Subject: [PATCH] Update export.py with v3.0 Hardswish() support (#831) --- models/export.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/models/export.py b/models/export.py index 00e3b2a4f..e870deca8 100644 --- a/models/export.py +++ b/models/export.py @@ -7,6 +7,7 @@ Usage: import argparse import torch +import torch.nn as nn from models.common import Conv from models.experimental import attempt_load @@ -32,7 +33,7 @@ if __name__ == '__main__': # Update model for k, m in model.named_modules(): m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability - if isinstance(m, Conv): + if isinstance(m, Conv) and isinstance(m.act, nn.Hardswish): m.act = Hardswish() # assign activation # if isinstance(m, Detect): # m.forward = m.forward_export # assign forward (optional)