Update davit.py

pull/1630/head
Fredo Guan 3 years ago
parent c9f42315de
commit 8792e906ab

@ -549,7 +549,7 @@ def _create_davit(variant, pretrained=False, **kwargs):
model_cls = DaViT model_cls = DaViT
features_only = False features_only = False
kwargs_filter = None kwargs_filter = None
if model_kwargs.pop('features_only', False): if kwargs.pop('features_only', False):
model_cls = DaViTFeatures model_cls = DaViTFeatures
kwargs_filter = ('num_classes', 'global_pool') kwargs_filter = ('num_classes', 'global_pool')
features_only = True features_only = True

Loading…
Cancel
Save