diff --git a/timm/models/davit.py b/timm/models/davit.py index 26a120fa..bde025cb 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -458,7 +458,7 @@ class DaViT(nn.Module): self.apply(self._init_weights) - self._update_forward_fn() + self.forward = self._get_forward_fn() ''' if self._features_only == True: self.forward = self.forward_features_full