Update backbone.py

datvuthanh · Oct 11, 2022 · 48ff4a0 · 48ff4a0
1 parent d53d2ea
commit 48ff4a0
Showing 1 changed file with 6 additions and 7 deletions.
diff --git a/backbone.py b/backbone.py
@@ -87,12 +87,11 @@ def __init__(self, num_classes=80, compound_coef=0, seg_classes=1, backbone_name
                 weights='imagenet',
             )
 
-        if not onnx_export:
-            self.anchors = Anchors(anchor_scale=self.anchor_scale[compound_coef],
-                                   pyramid_levels=(torch.arange(self.pyramid_levels[self.compound_coef]) + 3).tolist(),
-                                   onnx_export=onnx_export,
-                                   **kwargs)
-        else:
+        self.anchors = Anchors(anchor_scale=self.anchor_scale[compound_coef],
+                               pyramid_levels=(torch.arange(self.pyramid_levels[self.compound_coef]) + 3).tolist(),
+                               onnx_export=onnx_export,
+                               **kwargs)
+        if onnx_export:
             ## TODO: timm
             self.encoder.set_swish(memory_efficient=False)
 
@@ -121,9 +120,9 @@ def forward(self, inputs):
 
         regression = self.regressor(features)
         classification = self.classifier(features)
+        anchors = self.anchors(inputs, inputs.dtype)
 
         if not self.onnx_export:
-            anchors = self.anchors(inputs, inputs.dtype)
             return features, regression, classification, anchors, segmentation
         else:
             return regression, classification, segmentation