Update davit.py

3 years ago · 8f32133dbb
parent 3b993a9301
commit 8f32133dbb
1 changed files with 1 additions and 1 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -382,7 +382,7 @@ class DaViT(nn.Module):
        attn_drop_rate=0.,
        img_size=224,
        num_classes=1000,
-        global_pool='avg'
+        global_pool='avg',
        features_only = False
    ):
        super().__init__()