anatomy: ReleasedVNet.Config(
        channels_in=1,
        channels_out=4,
        channels_hidden=(16, 32, 64, 128, 256),
        input=VNetInput.Config(
            channels_in=1,
            channels_out=16,
            layer=ConvNorm.Config(
                channels_in=1,
                channels_out=16,
                conv=Conv2d.Config(
                    channels_in=1,
                    channels_out=16,
                    kernel_size=5,
                    stride=1,
                    padding=2,
                    dilation=1,
                    groups=1,
                    bias=True,
                    depth_index=(),
                    device=None,
                    dtype=None,
                    init_weight=priml.model.init.kaiming_uniform,
                    init_bias=torch.nn.init.zeros_
                ),
                norm=BatchNorm2d.Config(
                    channels_in=16,
                    channels_out=16,
                    momentum=0.1,
                    eps=1e-05,
                    elementwise_affine=True,
                    device=None,
                    dtype=None
                )
            ),
            activation=priml.model.swiglu.elu
        ),
        down=[
        │   VNetDown.Config(
        │       channels_in=16,
        │       channels_out=32,
        │       num_layers=1,
        │       proj=ConvNorm.Config(
        │           channels_in=16,
        │           channels_out=32,
        │           conv=Conv2d.Config(
        │               channels_in=16,
        │               channels_out=32,
        │               kernel_size=2,
        │               stride=2,
        │               padding=0,
        │               dilation=1,
        │               groups=1,
        │               bias=True,
        │               depth_index=(),
        │               device=None,
        │               dtype=None,
        │               init_weight=priml.model.init.kaiming_uniform,
        │               init_bias=torch.nn.init.zeros_
        │           ),
        │           norm=BatchNorm2d.Config(
        │               channels_in=32,
        │               channels_out=32,
        │               momentum=0.1,
        │               eps=1e-05,
        │               elementwise_affine=True,
        │               device=None,
        │               dtype=None
        │           )
        │       ),
        │       block=ConvNorm.Config(
        │           channels_in=32,
        │           channels_out=32,
        │           conv=Conv2d.Config(
        │               channels_in=32,
        │               channels_out=32,
        │               kernel_size=5,
        │               stride=1,
        │               padding=2,
        │               dilation=1,
        │               groups=1,
        │               bias=True,
        │               depth_index=(),
        │               device=None,
        │               dtype=None,
        │               init_weight=priml.model.init.kaiming_uniform,
        │               init_bias=torch.nn.init.zeros_
        │           ),
        │           norm=BatchNorm2d.Config(
        │               channels_in=32,
        │               channels_out=32,
        │               momentum=0.1,
        │               eps=1e-05,
        │               elementwise_affine=True,
        │               device=None,
        │               dtype=None
        │           )
        │       ),
        │       dropout=None,
        │       activation=priml.model.swiglu.elu
        │   ),
        │   VNetDown.Config(
        │       channels_in=32,
        │       channels_out=64,
        │       num_layers=2,
        │       proj=ConvNorm.Config(
        │           channels_in=32,
        │           channels_out=64,
        │           conv=Conv2d.Config(
        │               channels_in=32,
        │               channels_out=64,
        │               kernel_size=2,
        │               stride=2,
        │               padding=0,
        │               dilation=1,
        │               groups=1,
        │               bias=True,
        │               depth_index=(),
        │               device=None,
        │               dtype=None,
        │               init_weight=priml.model.init.kaiming_uniform,
        │               init_bias=torch.nn.init.zeros_
        │           ),
        │           norm=BatchNorm2d.Config(
        │               channels_in=64,
        │               channels_out=64,
        │               momentum=0.1,
        │               eps=1e-05,
        │               elementwise_affine=True,
        │               device=None,
        │               dtype=None
        │           )
        │       ),
        │       block=ConvNorm.Config(
        │           channels_in=64,
        │           channels_out=64,
        │           conv=Conv2d.Config(
        │               channels_in=64,
        │               channels_out=64,
        │               kernel_size=5,
        │               stride=1,
        │               padding=2,
        │               dilation=1,
        │               groups=1,
        │               bias=True,
        │               depth_index=(),
        │               device=None,
        │               dtype=None,
        │               init_weight=priml.model.init.kaiming_uniform,
        │               init_bias=torch.nn.init.zeros_
        │           ),
        │           norm=BatchNorm2d.Config(
        │               channels_in=64,
        │               channels_out=64,
        │               momentum=0.1,
        │               eps=1e-05,
        │               elementwise_affine=True,
        │               device=None,
        │               dtype=None
        │           )
        │       ),
        │       dropout=None,
        │       activation=priml.model.swiglu.elu
        │   ),
        │   VNetDown.Config(
        │       channels_in=64,
        │       channels_out=128,
        │       num_layers=3,
        │       proj=ConvNorm.Config(
        │           channels_in=64,
        │           channels_out=128,
        │           conv=Conv2d.Config(
        │               channels_in=64,
        │               channels_out=128,
        │               kernel_size=2,
        │               stride=2,
        │               padding=0,
        │               dilation=1,
        │               groups=1,
        │               bias=True,
        │               depth_index=(),
        │               device=None,
        │               dtype=None,
        │               init_weight=priml.model.init.kaiming_uniform,
        │               init_bias=torch.nn.init.zeros_
        │           ),
        │           norm=BatchNorm2d.Config(
        │               channels_in=128,
        │               channels_out=128,
        │               momentum=0.1,
        │               eps=1e-05,
        │               elementwise_affine=True,
        │               device=None,
        │               dtype=None
        │           )
        │       ),
        │       block=ConvNorm.Config(
        │           channels_in=128,
        │           channels_out=128,
        │           conv=Conv2d.Config(
        │               channels_in=128,
        │               channels_out=128,
        │               kernel_size=5,
        │               stride=1,
        │               padding=2,
        │               dilation=1,
        │               groups=1,
        │               bias=True,
        │               depth_index=(),
        │               device=None,
        │               dtype=None,
        │               init_weight=priml.model.init.kaiming_uniform,
        │               init_bias=torch.nn.init.zeros_
        │           ),
        │           norm=BatchNorm2d.Config(
        │               channels_in=128,
        │               channels_out=128,
        │               momentum=0.1,
        │               eps=1e-05,
        │               elementwise_affine=True,
        │               device=None,
        │               dtype=None
        │           )
        │       ),
        │       dropout=Dropout2d.Config(
        │           channels_in=128,
        │           channels_out=128,
        │           p=0.0,
        │           inplace=False
        │       ),
        │       activation=priml.model.swiglu.elu
        │   ),
        │   VNetDown.Config(
        │       channels_in=128,
        │       channels_out=256,
        │       num_layers=2,
        │       proj=ConvNorm.Config(
        │           channels_in=128,
        │           channels_out=256,
        │           conv=Conv2d.Config(
        │               channels_in=128,
        │               channels_out=256,
        │               kernel_size=2,
        │               stride=2,
        │               padding=0,
        │               dilation=1,
        │               groups=1,
        │               bias=True,
        │               depth_index=(),
        │               device=None,
        │               dtype=None,
        │               init_weight=priml.model.init.kaiming_uniform,
        │               init_bias=torch.nn.init.zeros_
        │           ),
        │           norm=BatchNorm2d.Config(
        │               channels_in=256,
        │               channels_out=256,
        │               momentum=0.1,
        │               eps=1e-05,
        │               elementwise_affine=True,
        │               device=None,
        │               dtype=None
        │           )
        │       ),
        │       block=ConvNorm.Config(
        │           channels_in=256,
        │           channels_out=256,
        │           conv=Conv2d.Config(
        │               channels_in=256,
        │               channels_out=256,
        │               kernel_size=5,
        │               stride=1,
        │               padding=2,
        │               dilation=1,
        │               groups=1,
        │               bias=True,
        │               depth_index=(),
        │               device=None,
        │               dtype=None,
        │               init_weight=priml.model.init.kaiming_uniform,
        │               init_bias=torch.nn.init.zeros_
        │           ),
        │           norm=BatchNorm2d.Config(
        │               channels_in=256,
        │               channels_out=256,
        │               momentum=0.1,
        │               eps=1e-05,
        │               elementwise_affine=True,
        │               device=None,
        │               dtype=None
        │           )
        │       ),
        │       dropout=Dropout2d.Config(
        │           channels_in=256,
        │           channels_out=256,
        │           p=0.0,
        │           inplace=False
        │       ),
        │       activation=priml.model.swiglu.elu
        │   )
        ],
        up=[
        │   VNetUp.Config(
        │       channels_in=256,
        │       channels_skip=128,
        │       num_layers=2,
        │       proj=ConvTranspose2d.Config(
        │           channels_in=256,
        │           channels_out=128,
        │           kernel_size=2,
        │           stride=2,
        │           padding=0,
        │           output_padding=0,
        │           dilation=1,
        │           groups=1,
        │           bias=True,
        │           depth_index=(),
        │           device=None,
        │           dtype=None,
        │           init_weight=priml.model.init.kaiming_uniform,
        │           init_bias=torch.nn.init.zeros_
        │       ),
        │       norm=BatchNorm2d.Config(
        │           channels_in=128,
        │           channels_out=128,
        │           momentum=0.1,
        │           eps=1e-05,
        │           elementwise_affine=True,
        │           device=None,
        │           dtype=None
        │       ),
        │       block=ConvNorm.Config(
        │           channels_in=256,
        │           channels_out=256,
        │           conv=Conv2d.Config(
        │               channels_in=256,
        │               channels_out=256,
        │               kernel_size=5,
        │               stride=1,
        │               padding=2,
        │               dilation=1,
        │               groups=1,
        │               bias=True,
        │               depth_index=(),
        │               device=None,
        │               dtype=None,
        │               init_weight=priml.model.init.kaiming_uniform,
        │               init_bias=torch.nn.init.zeros_
        │           ),
        │           norm=BatchNorm2d.Config(
        │               channels_in=256,
        │               channels_out=256,
        │               momentum=0.1,
        │               eps=1e-05,
        │               elementwise_affine=True,
        │               device=None,
        │               dtype=None
        │           )
        │       ),
        │       dropout=Dropout2d.Config(
        │           channels_in=256,
        │           channels_out=256,
        │           p=0.0,
        │           inplace=False
        │       ),
        │       dropout_skip=Dropout2d.Config(
        │           channels_in=128,
        │           channels_out=128,
        │           p=0.5,
        │           inplace=False
        │       ),
        │       activation=priml.model.swiglu.elu
        │   ),
        │   VNetUp.Config(
        │       channels_in=256,
        │       channels_skip=64,
        │       num_layers=2,
        │       proj=ConvTranspose2d.Config(
        │           channels_in=256,
        │           channels_out=64,
        │           kernel_size=2,
        │           stride=2,
        │           padding=0,
        │           output_padding=0,
        │           dilation=1,
        │           groups=1,
        │           bias=True,
        │           depth_index=(),
        │           device=None,
        │           dtype=None,
        │           init_weight=priml.model.init.kaiming_uniform,
        │           init_bias=torch.nn.init.zeros_
        │       ),
        │       norm=BatchNorm2d.Config(
        │           channels_in=64,
        │           channels_out=64,
        │           momentum=0.1,
        │           eps=1e-05,
        │           elementwise_affine=True,
        │           device=None,
        │           dtype=None
        │       ),
        │       block=ConvNorm.Config(
        │           channels_in=128,
        │           channels_out=128,
        │           conv=Conv2d.Config(
        │               channels_in=128,
        │               channels_out=128,
        │               kernel_size=5,
        │               stride=1,
        │               padding=2,
        │               dilation=1,
        │               groups=1,
        │               bias=True,
        │               depth_index=(),
        │               device=None,
        │               dtype=None,
        │               init_weight=priml.model.init.kaiming_uniform,
        │               init_bias=torch.nn.init.zeros_
        │           ),
        │           norm=BatchNorm2d.Config(
        │               channels_in=128,
        │               channels_out=128,
        │               momentum=0.1,
        │               eps=1e-05,
        │               elementwise_affine=True,
        │               device=None,
        │               dtype=None
        │           )
        │       ),
        │       dropout=Dropout2d.Config(
        │           channels_in=256,
        │           channels_out=256,
        │           p=0.0,
        │           inplace=False
        │       ),
        │       dropout_skip=Dropout2d.Config(
        │           channels_in=64,
        │           channels_out=64,
        │           p=0.5,
        │           inplace=False
        │       ),
        │       activation=priml.model.swiglu.elu
        │   ),
        │   VNetUp.Config(
        │       channels_in=128,
        │       channels_skip=32,
        │       num_layers=1,
        │       proj=ConvTranspose2d.Config(
        │           channels_in=128,
        │           channels_out=32,
        │           kernel_size=2,
        │           stride=2,
        │           padding=0,
        │           output_padding=0,
        │           dilation=1,
        │           groups=1,
        │           bias=True,
        │           depth_index=(),
        │           device=None,
        │           dtype=None,
        │           init_weight=priml.model.init.kaiming_uniform,
        │           init_bias=torch.nn.init.zeros_
        │       ),
        │       norm=BatchNorm2d.Config(
        │           channels_in=32,
        │           channels_out=32,
        │           momentum=0.1,
        │           eps=1e-05,
        │           elementwise_affine=True,
        │           device=None,
        │           dtype=None
        │       ),
        │       block=ConvNorm.Config(
        │           channels_in=64,
        │           channels_out=64,
        │           conv=Conv2d.Config(
        │               channels_in=64,
        │               channels_out=64,
        │               kernel_size=5,
        │               stride=1,
        │               padding=2,
        │               dilation=1,
        │               groups=1,
        │               bias=True,
        │               depth_index=(),
        │               device=None,
        │               dtype=None,
        │               init_weight=priml.model.init.kaiming_uniform,
        │               init_bias=torch.nn.init.zeros_
        │           ),
        │           norm=BatchNorm2d.Config(
        │               channels_in=64,
        │               channels_out=64,
        │               momentum=0.1,
        │               eps=1e-05,
        │               elementwise_affine=True,
        │               device=None,
        │               dtype=None
        │           )
        │       ),
        │       dropout=None,
        │       dropout_skip=Dropout2d.Config(
        │           channels_in=32,
        │           channels_out=32,
        │           p=0.5,
        │           inplace=False
        │       ),
        │       activation=priml.model.swiglu.elu
        │   ),
        │   VNetUp.Config(
        │       channels_in=64,
        │       channels_skip=16,
        │       num_layers=1,
        │       proj=ConvTranspose2d.Config(
        │           channels_in=64,
        │           channels_out=16,
        │           kernel_size=2,
        │           stride=2,
        │           padding=0,
        │           output_padding=0,
        │           dilation=1,
        │           groups=1,
        │           bias=True,
        │           depth_index=(),
        │           device=None,
        │           dtype=None,
        │           init_weight=priml.model.init.kaiming_uniform,
        │           init_bias=torch.nn.init.zeros_
        │       ),
        │       norm=BatchNorm2d.Config(
        │           channels_in=16,
        │           channels_out=16,
        │           momentum=0.1,
        │           eps=1e-05,
        │           elementwise_affine=True,
        │           device=None,
        │           dtype=None
        │       ),
        │       block=ConvNorm.Config(
        │           channels_in=32,
        │           channels_out=32,
        │           conv=Conv2d.Config(
        │               channels_in=32,
        │               channels_out=32,
        │               kernel_size=5,
        │               stride=1,
        │               padding=2,
        │               dilation=1,
        │               groups=1,
        │               bias=True,
        │               depth_index=(),
        │               device=None,
        │               dtype=None,
        │               init_weight=priml.model.init.kaiming_uniform,
        │               init_bias=torch.nn.init.zeros_
        │           ),
        │           norm=BatchNorm2d.Config(
        │               channels_in=32,
        │               channels_out=32,
        │               momentum=0.1,
        │               eps=1e-05,
        │               elementwise_affine=True,
        │               device=None,
        │               dtype=None
        │           )
        │       ),
        │       dropout=None,
        │       dropout_skip=Dropout2d.Config(
        │           channels_in=16,
        │           channels_out=16,
        │           p=0.5,
        │           inplace=False
        │       ),
        │       activation=priml.model.swiglu.elu
        │   )
        ],
        output=VNetOutput.Config(
            channels_in=32,
            channels_out=4,
            conv=ConvNorm.Config(
                channels_in=32,
                channels_out=4,
                conv=Conv2d.Config(
                    channels_in=32,
                    channels_out=4,
                    kernel_size=5,
                    stride=1,
                    padding=2,
                    dilation=1,
                    groups=1,
                    bias=True,
                    depth_index=(),
                    device=None,
                    dtype=None,
                    init_weight=priml.model.init.kaiming_uniform,
                    init_bias=torch.nn.init.zeros_
                ),
                norm=BatchNorm2d.Config(
                    channels_in=4,
                    channels_out=4,
                    momentum=0.1,
                    eps=1e-05,
                    elementwise_affine=True,
                    device=None,
                    dtype=None
                )
            ),
            proj=Conv2d.Config(
                channels_in=4,
                channels_out=4,
                kernel_size=1,
                stride=1,
                padding=0,
                dilation=1,
                groups=1,
                bias=True,
                depth_index=(),
                device=None,
                dtype=None,
                init_weight=priml.model.init.kaiming_uniform,
                init_bias=torch.nn.init.zeros_
            ),
            activation=priml.model.swiglu.elu
        )
    )
side: ResNet.Config(
        channels_in=1,
        channels_out=1,
        stem=Conv2d.Config(
            channels_in=1,
            channels_out=64,
            kernel_size=7,
            stride=2,
            padding=3,
            dilation=1,
            groups=1,
            bias=False,
            depth_index=(),
            device=None,
            dtype=None,
            init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
            init_bias=torch.nn.init.zeros_
        ),
        norm=BatchNorm2d.Config(
            channels_in=64,
            channels_out=64,
            momentum=0.1,
            eps=1e-05,
            elementwise_affine=True,
            device=None,
            dtype=None
        ),
        channels_hidden=(64, 128, 256, 512),
        blocks_per_stage=(3, 4, 6, 3),
        stage=[
        │   ResNetStage.Config(
        │       channels_in=64,
        │       channels_out=64,
        │       stride=1,
        │       blocks_per_stage=3,
        │       block=[
        │       │   ResNetBlock.Config(
        │       │       channels_in=64,
        │       │       channels_out=64,
        │       │       stride=1,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           kernel_size=1,
        │       │           stride=1,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   ),
        │       │   ResNetBlock.Config(
        │       │       channels_in=64,
        │       │       channels_out=64,
        │       │       stride=1,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           kernel_size=1,
        │       │           stride=1,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   ),
        │       │   ResNetBlock.Config(
        │       │       channels_in=64,
        │       │       channels_out=64,
        │       │       stride=1,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           kernel_size=1,
        │       │           stride=1,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=64,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   )
        │       ]
        │   ),
        │   ResNetStage.Config(
        │       channels_in=64,
        │       channels_out=128,
        │       stride=2,
        │       blocks_per_stage=4,
        │       block=[
        │       │   ResNetBlock.Config(
        │       │       channels_in=64,
        │       │       channels_out=128,
        │       │       stride=2,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=128,
        │       │           kernel_size=3,
        │       │           stride=2,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=64,
        │       │           channels_out=128,
        │       │           kernel_size=1,
        │       │           stride=2,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   ),
        │       │   ResNetBlock.Config(
        │       │       channels_in=128,
        │       │       channels_out=128,
        │       │       stride=1,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           kernel_size=1,
        │       │           stride=1,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   ),
        │       │   ResNetBlock.Config(
        │       │       channels_in=128,
        │       │       channels_out=128,
        │       │       stride=1,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           kernel_size=1,
        │       │           stride=1,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   ),
        │       │   ResNetBlock.Config(
        │       │       channels_in=128,
        │       │       channels_out=128,
        │       │       stride=1,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           kernel_size=1,
        │       │           stride=1,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=128,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   )
        │       ]
        │   ),
        │   ResNetStage.Config(
        │       channels_in=128,
        │       channels_out=256,
        │       stride=2,
        │       blocks_per_stage=6,
        │       block=[
        │       │   ResNetBlock.Config(
        │       │       channels_in=128,
        │       │       channels_out=256,
        │       │       stride=2,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=256,
        │       │           kernel_size=3,
        │       │           stride=2,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=128,
        │       │           channels_out=256,
        │       │           kernel_size=1,
        │       │           stride=2,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   ),
        │       │   ResNetBlock.Config(
        │       │       channels_in=256,
        │       │       channels_out=256,
        │       │       stride=1,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=1,
        │       │           stride=1,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   ),
        │       │   ResNetBlock.Config(
        │       │       channels_in=256,
        │       │       channels_out=256,
        │       │       stride=1,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=1,
        │       │           stride=1,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   ),
        │       │   ResNetBlock.Config(
        │       │       channels_in=256,
        │       │       channels_out=256,
        │       │       stride=1,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=1,
        │       │           stride=1,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   ),
        │       │   ResNetBlock.Config(
        │       │       channels_in=256,
        │       │       channels_out=256,
        │       │       stride=1,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=1,
        │       │           stride=1,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   ),
        │       │   ResNetBlock.Config(
        │       │       channels_in=256,
        │       │       channels_out=256,
        │       │       stride=1,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           kernel_size=1,
        │       │           stride=1,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=256,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   )
        │       ]
        │   ),
        │   ResNetStage.Config(
        │       channels_in=256,
        │       channels_out=512,
        │       stride=2,
        │       blocks_per_stage=3,
        │       block=[
        │       │   ResNetBlock.Config(
        │       │       channels_in=256,
        │       │       channels_out=512,
        │       │       stride=2,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=512,
        │       │           kernel_size=3,
        │       │           stride=2,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=256,
        │       │           channels_out=512,
        │       │           kernel_size=1,
        │       │           stride=2,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   ),
        │       │   ResNetBlock.Config(
        │       │       channels_in=512,
        │       │       channels_out=512,
        │       │       stride=1,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           kernel_size=1,
        │       │           stride=1,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   ),
        │       │   ResNetBlock.Config(
        │       │       channels_in=512,
        │       │       channels_out=512,
        │       │       stride=1,
        │       │       conv1=Conv2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm1=BatchNorm2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       conv2=Conv2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           kernel_size=3,
        │       │           stride=1,
        │       │           padding=1,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm2=BatchNorm2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       projection=Conv2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           kernel_size=1,
        │       │           stride=1,
        │       │           padding=0,
        │       │           dilation=1,
        │       │           groups=1,
        │       │           bias=False,
        │       │           depth_index=(),
        │       │           device=None,
        │       │           dtype=None,
        │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │       │           init_bias=torch.nn.init.zeros_
        │       │       ),
        │       │       norm_projection=BatchNorm2d.Config(
        │       │           channels_in=512,
        │       │           channels_out=512,
        │       │           momentum=0.1,
        │       │           eps=1e-05,
        │       │           elementwise_affine=True,
        │       │           device=None,
        │       │           dtype=None
        │       │       ),
        │       │       activation=priml.model.resnet.relu_inplace,
        │       │       init_scale_residual=torch.nn.init.zeros_
        │       │   )
        │       ]
        │   )
        ],
        proj_out=Linear.Config(
            channels_in=512,
            channels_out=1,
            bias=True,
            depth_index=(),
            device=None,
            dtype=None,
            init_weight=priml.model.init.kaiming_uniform,
            init_bias=torch.nn.init.zeros_,
            linear_fn=<built-in function linear>,
            shard=None
        ),
        activation=priml.model.resnet.relu_inplace
    )
