fix fx tracer to support alpha in RepVGGBlock

Signed-off-by: alec.tu <[email protected]>
open-mmlab · Aug 29, 2023 · a2df191 · a2df191
1 parent 106bec4
commit a2df191
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/mmyolo/models/layers/yolo_bricks.py b/mmyolo/models/layers/yolo_bricks.py
@@ -183,6 +183,7 @@ def __init__(self,
         else:
             self.se = nn.Identity()
 
+        self.use_alpha = use_alpha
         if use_alpha:
             alpha = torch.ones([
                 1,
@@ -246,7 +247,7 @@ def forward(self, inputs: Tensor) -> Tensor:
             id_out = 0
         else:
             id_out = self.rbr_identity(inputs)
-        if self.alpha:
+        if self.use_alpha:
             return self.nonlinearity(
                 self.se(
                     self.rbr_dense(inputs) +
@@ -265,7 +266,7 @@ def get_equivalent_kernel_bias(self):
         kernel3x3, bias3x3 = self._fuse_bn_tensor(self.rbr_dense)
         kernel1x1, bias1x1 = self._fuse_bn_tensor(self.rbr_1x1)
         kernelid, biasid = self._fuse_bn_tensor(self.rbr_identity)
-        if self.alpha:
+        if self.use_alpha:
             return kernel3x3 + self.alpha * self._pad_1x1_to_3x3_tensor(
                 kernel1x1) + kernelid, bias3x3 + self.alpha * bias1x1 + biasid
         else: