Fix last min torchscript regression in nfnet changes

rwightman · rwightman · commit d196fa536dca · 2023-03-24T00:10:17.000-07:00
diff --git a/timm/models/nfnet.py b/timm/models/nfnet.py
@@ -445,7 +445,7 @@ def forward_features(self, x):
         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):
