Add global_pool to mambaout __init__ and pass to heads

Author: rwightman

timm/models/mambaout.py

@@ -300,6 +300,7 @@ def __init__(
             self,
             in_chans=3,
             num_classes=1000,
+            global_pool='avg',
             depths=(3, 3, 9, 3),
             dims=(96, 192, 384, 576),
             norm_layer=LayerNorm,
@@ -369,7 +370,7 @@ def __init__(
             self.head = MlpHead(
                 prev_dim,
                 num_classes,
-                pool_type='avg',
+                pool_type=global_pool,
                 drop_rate=drop_rate,
                 norm_layer=norm_layer,
             )
@@ -379,7 +380,7 @@ def __init__(
                 prev_dim,
                 num_classes,
                 hidden_size=int(prev_dim * 4),
-                pool_type='avg',
+                pool_type=global_pool,
                 norm_layer=norm_layer,
                 drop_rate=drop_rate,
             )