set self.num_features to neck_chans if neck_chans > 0

huggingface · Aug 11, 2023 · 2048f6f · 2048f6f
1 parent 3a44e6c
commit 2048f6f
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/timm/models/vision_transformer_sam.py b/timm/models/vision_transformer_sam.py
@@ -434,6 +434,7 @@ def __init__(
                 ),
                 LayerNorm2d(neck_chans),
             )
+            self.num_features = neck_chans
         else:
             self.neck = nn.Identity()
             neck_chans = embed_dim