See if we can avoid some model / layer pickle issues with the aa attr in ConvNormAct

2024-12-02 16:55:29 -08:00 · 2024-12-02 16:55:29 -08:00 · 464885e135
parent 5fe5f9d488
commit 464885e135
1 changed files with 2 additions and 1 deletions
--- a/timm/layers/conv_bn_act.py
+++ b/timm/layers/conv_bn_act.py
@ -82,7 +82,8 @@ class ConvNormAct(nn.Module):
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
-        if self.aa is not None:
+        aa = getattr(self, 'aa', None)
+        if aa is not None:
            x = self.aa(x)
        return x