From 2f77cf33f651d7c6b46abf696cd7020022c13982 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 20 Sep 2020 12:04:20 -0700 Subject: [PATCH] .fuse() additional error checking --- models/yolo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/yolo.py b/models/yolo.py index f9929aa..4b2606d 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -160,7 +160,7 @@ class Model(nn.Module): def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers print('Fusing layers... ') for m in self.model.modules(): - if type(m) is Conv: + if type(m) is Conv and hasattr(Conv, 'bn'): m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv delattr(m, 'bn') # remove batchnorm