Update export.py with v3.0 Hardswish() support (#831)

ultralytics · Aug 25, 2020 · 4fb8cb3 · 4fb8cb3
1 parent 4d7f222
commit 4fb8cb3
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/models/export.py b/models/export.py
@@ -7,6 +7,7 @@
 import argparse
 
 import torch
+import torch.nn as nn
 
 from models.common import Conv
 from models.experimental import attempt_load
@@ -32,7 +33,7 @@
     # 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)