From 0ddadad723e4d9d97884079468fa4c9b904faaff Mon Sep 17 00:00:00 2001
From: caitianren <295727207@qq.com>
Date: Wed, 29 Nov 2023 20:15:00 +0800
Subject: [PATCH] Fix this bug when setting distillation-type to none

---
 models/swiftformer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models/swiftformer.py b/models/swiftformer.py
index 1b74936..b545557 100644
--- a/models/swiftformer.py
+++ b/models/swiftformer.py
@@ -437,7 +437,7 @@ class SwiftFormer(nn.Module):
             if not self.training:
                 cls_out = (cls_out[0] + cls_out[1]) / 2
         else:
-            cls_out = self.head(x.mean(-2))
+            cls_out = self.head(x.flatten(2).mean(-1))
         # For image classification
         return cls_out