Skip to content

Commit

Permalink
Add missing dropout for head reset in ConvNeXt default head
Browse files Browse the repository at this point in the history
  • Loading branch information
rwightman committed Jan 15, 2022
1 parent b093dcb commit edd3d73
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion timm/models/convnext.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
""" ConvNext
""" ConvNeXt
Paper: `A ConvNet for the 2020s` - https://arxiv.org/pdf/2201.03545.pdf
Expand Down Expand Up @@ -229,6 +229,7 @@ def __init__(
('global_pool', SelectAdaptivePool2d(pool_type=global_pool)),
('norm', norm_layer(self.num_features)),
('flatten', nn.Flatten(1) if global_pool else nn.Identity()),
('drop', nn.Dropout(self.drop_rate)),
('fc', nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity())
]))

Expand Down

0 comments on commit edd3d73

Please sign in to comment.