filter kwargs before creating norm layer

timm/layers/create_norm.py

@@ -5,6 +5,7 @@
 Copyright 2022 Ross Wightman
 """
 import functools
+import inspect
 import types
 from typing import Type
 
@@ -48,12 +49,14 @@
     simplenorm2dfp32=SimpleNorm2dFp32,
     frozenbatchnorm2d=FrozenBatchNorm2d,
 )
-_NORM_TYPES = {m for n, m in _NORM_MAP.items()}
 
 
-def create_norm_layer(layer_name, num_features, **kwargs):
+def create_norm_layer(layer_name, **kwargs):
     layer = get_norm_layer(layer_name)
-    layer_instance = layer(num_features, **kwargs)
+    # filter kwargs before creating layer
+    params = inspect.signature(layer).parameters
+    kwargs = {k: v for k, v in kwargs.items() if k in params}
+    layer_instance = layer(**kwargs)
     return layer_instance
 
 

timm/models/efficientformer_v2.py

@@ -96,7 +96,14 @@ def __init__(
             bias=bias,
             **dd,
         )
-        self.bn = create_norm_layer(norm_layer, out_channels, **norm_kwargs, **dd)
+        self.bn = create_norm_layer(
+            norm_layer,
+            num_channels=out_channels,
+            num_features=out_channels,
+            num_groups=groups,
+            **norm_kwargs,
+            **dd,
+        )
 
     def forward(self, x):
         x = self.conv(x)