typing MC

Author: ayasyrev

model_constructor/model_constructor.py

@@ -150,22 +150,22 @@ def __init__(
         expansion: int = 1,
         groups: int = 1,
         dw: bool = False,
-        div_groups=None,
+        div_groups: Union[int, None]=None,
         sa: Union[bool, int, Callable] = False,
         se: Union[bool, int, Callable] = False,
         se_module=None,
         se_reduction=None,
-        bn_1st=True,
-        zero_bn=True,
-        stem_stride_on=0,
-        stem_sizes=[32, 32, 64],
-        stem_pool=nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
-        stem_bn_end=False,
-        _init_cnn=init_cnn,
-        _make_stem=_make_stem,
-        _make_layer=_make_layer,
-        _make_body=_make_body,
-        _make_head=_make_head,
+        bn_1st: bool = True,
+        zero_bn: bool = True,
+        stem_stride_on: int = 0,
+        stem_sizes: List[int] = [32, 32, 64],
+        stem_pool: Union[nn.Module, None] =nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
+        stem_bn_end: bool = False,
+        _init_cnn: Callable = init_cnn,
+        _make_stem: Callable = _make_stem,
+        _make_layer: Callable = _make_layer,
+        _make_body: Callable = _make_body,
+        _make_head: Callable = _make_head,
     ):
         super().__init__()
         # se can be bool, int (0, 1) or nn.Module