kwargs to mc create model

Author: ayasyrev

src/model_constructor/model_constructor.py

@@ -412,10 +412,10 @@ def from_cfg(cls, cfg: ModelCfg):
         return cls(**cfg.dict())
 
     @classmethod
-    def create_model(cls, cfg: Union[ModelCfg, None] = None) -> nn.Sequential:
+    def create_model(cls, cfg: Union[ModelCfg, None] = None, **kwargs: dict[str, Any]) -> nn.Sequential:
         if cfg:
             return cls(**cfg.dict())()
-        return cls()()
+        return cls(**kwargs)()
 
     def __call__(self) -> nn.Sequential:
         """Create model."""

tests/test_mc.py

@@ -56,6 +56,9 @@ def test_MC():
     pred = model(xb)
     assert pred.shape == torch.Size([bs_test, 1000])
 
+    model = ModelConstructor.create_model(num_classes=num_classes)
+    pred = model(xb)
+    assert pred.shape == torch.Size([bs_test, num_classes])
 
 
 def test_MC_bottleneck():