Make sure num_features attr is present in vit models as with others

pull/263/head
Ross Wightman 4 years ago
parent 4a3df7842a
commit f591e90b0d

@ -206,7 +206,7 @@ class VisionTransformer(nn.Module):
drop_path_rate=0., hybrid_backbone=None, norm_layer=nn.LayerNorm): drop_path_rate=0., hybrid_backbone=None, norm_layer=nn.LayerNorm):
super().__init__() super().__init__()
self.num_classes = num_classes self.num_classes = num_classes
self.embed_dim = embed_dim self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
if hybrid_backbone is not None: if hybrid_backbone is not None:
self.patch_embed = HybridEmbed( self.patch_embed = HybridEmbed(

Loading…
Cancel
Save