fix fx tracer to support alpha in RepVGGBlock

open-mmlab · Aug 29, 2023 · 3c61e21 · 3c61e21
1 parent 106bec4
commit 3c61e21
Showing 1 changed file with 2 additions and 1 deletion.
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) +