Stricter check on pool_type for adaptive pooling module. Fix #2159

Author: rwightman

timm/layers/adaptive_avgmax_pool.py

@@ -134,6 +134,7 @@ def __init__(
         super(SelectAdaptivePool2d, self).__init__()
         assert input_fmt in ('NCHW', 'NHWC')
         self.pool_type = pool_type or ''  # convert other falsy values to empty string for consistent TS typing
+        pool_type = pool_type.lower()
         if not pool_type:
             self.pool = nn.Identity()  # pass through
             self.flatten = nn.Flatten(1) if flatten else nn.Identity()
@@ -145,8 +146,10 @@ def __init__(
                 self.pool = FastAdaptiveAvgMaxPool(flatten, input_fmt=input_fmt)
             elif pool_type.endswith('max'):
                 self.pool = FastAdaptiveMaxPool(flatten, input_fmt=input_fmt)
-            else:
+            elif pool_type == 'fast' or pool_type.endswith('avg'):
                 self.pool = FastAdaptiveAvgPool(flatten, input_fmt=input_fmt)
+            else:
+                assert False, 'Invalid pool type: %s' % pool_type
             self.flatten = nn.Identity()
         else:
             assert input_fmt == 'NCHW'
@@ -156,8 +159,10 @@ def __init__(
                 self.pool = AdaptiveCatAvgMaxPool2d(output_size)
             elif pool_type == 'max':
                 self.pool = nn.AdaptiveMaxPool2d(output_size)
-            else:
+            elif pool_type == 'avg':
                 self.pool = nn.AdaptiveAvgPool2d(output_size)
+            else:
+                assert False, 'Invalid pool type: %s' % pool_type
             self.flatten = nn.Flatten(1) if flatten else nn.Identity()
 
     def is_identity(self):