diff --git a/mcunet/tinynas/elastic_nn/networks/ofa_proxyless.py b/mcunet/tinynas/elastic_nn/networks/ofa_proxyless.py
index a64cc3e..743161b 100644
--- a/mcunet/tinynas/elastic_nn/networks/ofa_proxyless.py
+++ b/mcunet/tinynas/elastic_nn/networks/ofa_proxyless.py
@@ -359,7 +359,7 @@ def get_or_copy_subnet(m, **kwargs):
                 input_channel = stage_blocks[-1].mobile_inverted_conv.out_channels
             blocks += stage_blocks
 
-        feature_mix_layer = get_or_copy_subnet(self.feature_mix_layer, input_channel=input_channel)
+        feature_mix_layer = get_or_copy_subnet(self.feature_mix_layer, in_channel=input_channel)
         input_channel = feature_mix_layer.out_channels if feature_mix_layer is not None else input_channel
         classifier = get_or_copy_subnet(self.classifier, in_features=input_channel)