Merge pull request #1903 from twmht/fix_num_classes

fix num_classes not found in repvit
huggingface · Aug 7, 2023 · f677190 · f677190
2 parents 81089b1 + bb2b6b5
commit f677190
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/timm/models/repvit.py b/timm/models/repvit.py
@@ -176,6 +176,7 @@ def __init__(self, dim, num_classes, distillation=False):
         super().__init__()
         self.head = NormLinear(dim, num_classes) if num_classes > 0 else nn.Identity()
         self.distillation = distillation
+        self.num_classes=num_classes
         if distillation:
             self.head_dist = NormLinear(dim, num_classes) if num_classes > 0 else nn.Identity()