update vitamin.py
Browse files- vitamin.py +1 -1
vitamin.py
CHANGED
@@ -765,7 +765,7 @@ def vitamin_large_384(pretrained=False, **kwargs) -> VisionTransformer:
|
|
765 |
head_type='1d',
|
766 |
),
|
767 |
)
|
768 |
-
model_args = dict(img_size=384, embed_dim=1024, depth=31, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False,
|
769 |
model = _create_vision_transformer_hybrid(
|
770 |
'vitamin_large_384', backbone=backbone, pretrained=pretrained, **dict(model_args, **kwargs))
|
771 |
return model
|
|
|
765 |
head_type='1d',
|
766 |
),
|
767 |
)
|
768 |
+
model_args = dict(img_size=384, embed_dim=1024, depth=31, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
|
769 |
model = _create_vision_transformer_hybrid(
|
770 |
'vitamin_large_384', backbone=backbone, pretrained=pretrained, **dict(model_args, **kwargs))
|
771 |
return model
|