From bb2b6b5f09737059ccca004c35f3e2e611d341db Mon Sep 17 00:00:00 2001 From: "alec.tu" Date: Mon, 7 Aug 2023 15:16:03 +0800 Subject: [PATCH] fix num_classes not found --- timm/models/repvit.py | 1 + 1 file changed, 1 insertion(+) diff --git a/timm/models/repvit.py b/timm/models/repvit.py index b0199b89..43e35be9 100644 --- a/timm/models/repvit.py +++ b/timm/models/repvit.py @@ -176,6 +176,7 @@ class RepViTClassifier(nn.Module): 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()