diff --git a/timm/models/davit.py b/timm/models/davit.py index ff69dc89..e32aedf8 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -519,8 +519,10 @@ class DaViT(nn.Module): return x def forward(self, x): - return self.forward_classifier(self, x) - + x = self.forward_classifier(self, x) + return x + + class DaViTFeatures(DaViT): def __init__(*args): @@ -528,8 +530,8 @@ class DaViTFeatures(DaViT): self.feature_info = FeatureInfo(self.feature_info, kwargs.get('out_inices', (1, 2, 3, 4))) def forward(self, x) -> List[Tensor]: - return self.forward_pyramid_features(self, x) - + x = self.forward_pyramid_features(self, x) + return x def checkpoint_filter_fn(state_dict, model):