Update backbonev2.py

snap-research · Jan 22, 2023 · 2c0e950 · 2c0e950
1 parent 5f1dddc
commit 2c0e950
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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