create model at MC

Author: ayasyrev

src/model_constructor/model_constructor.py

@@ -411,6 +411,12 @@ def body(self):
     def from_cfg(cls, cfg: ModelCfg):
         return cls(**cfg.dict())
 
+    @classmethod
+    def create_model(cls, cfg: Union[ModelCfg, None] = None) -> nn.Sequential:
+        if cfg:
+            return cls(**cfg.dict())()
+        return cls()()
+
     def __call__(self) -> nn.Sequential:
         """Create model."""
         model_name = self.name or self.__class__.__name__

tests/test_mc.py

@@ -51,6 +51,11 @@ def test_MC():
     model = mc()
     pred = model(xb)
     assert pred.shape == torch.Size([bs_test, num_classes])
+    # create model from class. Default config, num_classes 1000. ??- how to change
+    model = ModelConstructor.create_model()
+    pred = model(xb)
+    assert pred.shape == torch.Size([bs_test, 1000])
+
 
 
 def test_MC_bottleneck():