From 2c0e950dc269b7f0229a3917fd54661b964554e0 Mon Sep 17 00:00:00 2001
From: Yanyu <liyy201912@gmail.com>
Date: Sun, 22 Jan 2023 14:51:38 -0500
Subject: [PATCH] Update backbonev2.py

---
 segmentation/backbonev2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/segmentation/backbonev2.py b/segmentation/backbonev2.py
index fa00167..d4a91cc 100644
--- a/segmentation/backbonev2.py
+++ b/segmentation/backbonev2.py
@@ -668,7 +668,7 @@ def forward(self, x):
             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