Fix last min torchscript regression in nfnet changes

2023-03-24 00:10:17 -07:00 · 2023-03-24 00:10:17 -07:00 · d196fa536d
parent 33ada0cbca
commit d196fa536d
1 changed files with 1 additions and 1 deletions
--- a/timm/models/nfnet.py
+++ b/timm/models/nfnet.py
@ -445,7 +445,7 @@ class NormFreeNet(nn.Module):
        x = self.final_act(x)
        return x

-    def forward_head(self, x, pre_logits=False):
+    def forward_head(self, x, pre_logits: bool = False):
        return self.head(x, pre_logits=pre_logits) if pre_logits else self.head(x)

    def forward(self, x):