diff --git a/timm/models/davit.py b/timm/models/davit.py index 06f65020..d607f945 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -577,9 +577,9 @@ class DaViT(nn.Module): def forward(self, x): if self.features_only == True: - return forward_features_full(x) + return self.forward_features_full(x) else: - return forward(x) + return self.forward_classification(x) def checkpoint_filter_fn(state_dict, model): """ Remap MSFT checkpoints -> timm """