xresnet26

Author: ayasyrev

src/model_constructor/model_constructor.py

@@ -32,7 +32,7 @@
 }
 
 
-nnModule = Union[Type[nn.Module], Callable[[], nn.Module]]
+nnModule = Union[Type[nn.Module], Callable[[Any], nn.Module]]
 
 
 class ModelCfg(Cfg, arbitrary_types_allowed=True, extra="forbid"):

src/model_constructor/xresnet.py

@@ -50,6 +50,11 @@ class XResNet34(XResNet):
     layers: List[int] = [3, 4, 6, 3]
 
 
-class XResNet50(XResNet34):
+class XResNet26(XResNet):
     block: Type[nn.Module] = BottleneckBlock
     block_sizes: List[int] = [256, 512, 1024, 2048]
+    expansion: int = 4
+
+
+class XResNet50(XResNet26):
+    layers: List[int] = [3, 4, 6, 3]