raptor224: RaptorStage.Config(
        base_dir='/opt/scratch',
        checkpoint=ReleasedState.Config(
            base_dir='/opt/scratch',
            file=LocalFile.Config(
                path=PosixPath('/opt/scratch/datasets/rsna-knee-raptor-swa-v1/raptor_ft_alldata_t16_blendjev_oai_d96_c08_r224_swa.pt'),
                sha256='3394fdd7e668286956e85af2554bef0386df5cd79d1c1af223b095c828b815d6'
            ),
            key='model',
            weights_only=False
        ),
        model=Raptor.Config(
        │   backbone=CoAtNet.Config(
        │   │   image_size=(224, 224),
        │   │   stem=CoAtStem.Config(
        │   │       channels_in=3,
        │   │       channels_hidden=(64, 128),
        │   │       conv1=Conv2d.Config(
        │   │           channels_in=3,
        │   │           channels_out=64,
        │   │           kernel_size=3,
        │   │           stride=2,
        │   │           padding=1,
        │   │           dilation=1,
        │   │           groups=1,
        │   │           bias=False,
        │   │           depth_index=(),
        │   │           device=None,
        │   │           dtype=None,
        │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │           init_bias=torch.nn.init.zeros_
        │   │       ),
        │   │       norm1=CoAtNorm.Config(
        │   │           channels_in=64,
        │   │           channels_out=64,
        │   │           momentum=0.1,
        │   │           eps=1e-05,
        │   │           elementwise_affine=True,
        │   │           device=None,
        │   │           dtype=None,
        │   │           activation=priml.model.coatnet.silu_inplace
        │   │       ),
        │   │       conv2=Conv2d.Config(
        │   │           channels_in=64,
        │   │           channels_out=128,
        │   │           kernel_size=3,
        │   │           stride=1,
        │   │           padding=1,
        │   │           dilation=1,
        │   │           groups=1,
        │   │           bias=False,
        │   │           depth_index=(),
        │   │           device=None,
        │   │           dtype=None,
        │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │           init_bias=torch.nn.init.zeros_
        │   │       )
        │   │   ),
        │   │   channels_hidden=(128, 256, 512, 1_024),
        │   │   blocks_per_stage=(2, 6, 14, 2),
        │   │   stage=[
        │   │   │   CoAtStage.Config(
        │   │   │       channels_in=128,
        │   │   │       channels_out=128,
        │   │   │       feature_size=(56, 56),
        │   │   │       blocks_per_stage=2,
        │   │   │       transition=CoAtMBConv.Config(
        │   │   │       │   channels_in=128,
        │   │   │       │   channels_out=128,
        │   │   │       │   stride=2,
        │   │   │       │   expansion=4.0,
        │   │   │       │   pre_norm=CoAtNorm.Config(
        │   │   │       │       channels_in=128,
        │   │   │       │       channels_out=128,
        │   │   │       │       momentum=0.1,
        │   │   │       │       eps=1e-05,
        │   │   │       │       elementwise_affine=True,
        │   │   │       │       device=None,
        │   │   │       │       dtype=None,
        │   │   │       │       activation=priml.model.coatnet.silu_inplace
        │   │   │       │   ),
        │   │   │       │   conv1_1x1=Conv2d.Config(
        │   │   │       │       channels_in=128,
        │   │   │       │       channels_out=512,
        │   │   │       │       kernel_size=1,
        │   │   │       │       stride=1,
        │   │   │       │       padding=0,
        │   │   │       │       dilation=1,
        │   │   │       │       groups=1,
        │   │   │       │       bias=False,
        │   │   │       │       depth_index=(),
        │   │   │       │       device=None,
        │   │   │       │       dtype=None,
        │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │       init_bias=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   norm1=CoAtNorm.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       momentum=0.1,
        │   │   │       │       eps=1e-05,
        │   │   │       │       elementwise_affine=True,
        │   │   │       │       device=None,
        │   │   │       │       dtype=None,
        │   │   │       │       activation=priml.model.coatnet.silu_inplace
        │   │   │       │   ),
        │   │   │       │   conv2_kxk=Conv2d.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       kernel_size=3,
        │   │   │       │       stride=2,
        │   │   │       │       padding=1,
        │   │   │       │       dilation=1,
        │   │   │       │       groups=512,
        │   │   │       │       bias=False,
        │   │   │       │       depth_index=(),
        │   │   │       │       device=None,
        │   │   │       │       dtype=None,
        │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │       init_bias=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   norm2=CoAtNorm.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       momentum=0.1,
        │   │   │       │       eps=1e-05,
        │   │   │       │       elementwise_affine=True,
        │   │   │       │       device=None,
        │   │   │       │       dtype=None,
        │   │   │       │       activation=priml.model.coatnet.silu_inplace
        │   │   │       │   ),
        │   │   │       │   se=SqueezeExcitation.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       channels_hidden=128,
        │   │   │       │       fc1=Conv2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=128,
        │   │   │       │           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_
        │   │   │       │       ),
        │   │   │       │       fc2=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=512,
        │   │   │       │           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_
        │   │   │       │       ),
        │   │   │       │       squeeze=functools.partial(<built-in method mean of type object at 0xdefacedeface>, dim=(2, 3), keepdim=True),
        │   │   │       │       expand=None,
        │   │   │       │       act=priml.model.coatnet.silu_inplace,
        │   │   │       │       gate=priml.model.swiglu.sigmoid
        │   │   │       │   ),
        │   │   │       │   conv3_1x1=Conv2d.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=128,
        │   │   │       │       kernel_size=1,
        │   │   │       │       stride=1,
        │   │   │       │       padding=0,
        │   │   │       │       dilation=1,
        │   │   │       │       groups=1,
        │   │   │       │       bias=False,
        │   │   │       │       depth_index=(),
        │   │   │       │       device=None,
        │   │   │       │       dtype=None,
        │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │       init_bias=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   shortcut=CoAtPool.Config(
        │   │   │       │       channels_in=128,
        │   │   │       │       channels_out=128,
        │   │   │       │       pool=AvgPool2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=2,
        │   │   │       │           stride=2,
        │   │   │       │           padding=0,
        │   │   │       │           ceil_mode=False,
        │   │   │       │           count_include_pad=True,
        │   │   │       │           divisor_override=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=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       )
        │   │   │       │   )
        │   │   │       ),
        │   │   │       block=[
        │   │   │       │   CoAtMBConv.Config(
        │   │   │       │       channels_in=128,
        │   │   │       │       channels_out=128,
        │   │   │       │       stride=1,
        │   │   │       │       expansion=4.0,
        │   │   │       │       pre_norm=CoAtNorm.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       conv1_1x1=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=512,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=CoAtNorm.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       conv2_kxk=Conv2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=512,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=CoAtNorm.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       se=SqueezeExcitation.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           channels_hidden=128,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=128,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=128,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           squeeze=functools.partial(<built-in method mean of type object at 0xdefacedeface>, dim=(2, 3), keepdim=True),
        │   │   │       │           expand=None,
        │   │   │       │           act=priml.model.coatnet.silu_inplace,
        │   │   │       │           gate=priml.model.swiglu.sigmoid
        │   │   │       │       ),
        │   │   │       │       conv3_1x1=Conv2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       shortcut=CoAtPool.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           pool=AvgPool2d.Config(
        │   │   │       │               channels_in=128,
        │   │   │       │               channels_out=128,
        │   │   │       │               kernel_size=2,
        │   │   │       │               stride=2,
        │   │   │       │               padding=0,
        │   │   │       │               ceil_mode=False,
        │   │   │       │               count_include_pad=True,
        │   │   │       │               divisor_override=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=priml.model.init.kaiming_uniform,
        │   │   │       │               init_bias=torch.nn.init.zeros_
        │   │   │       │           )
        │   │   │       │       )
        │   │   │       │   )
        │   │   │       ]
        │   │   │   ),
        │   │   │   CoAtStage.Config(
        │   │   │       channels_in=128,
        │   │   │       channels_out=256,
        │   │   │       feature_size=(28, 28),
        │   │   │       blocks_per_stage=6,
        │   │   │       transition=CoAtMBConv.Config(
        │   │   │       │   channels_in=128,
        │   │   │       │   channels_out=256,
        │   │   │       │   stride=2,
        │   │   │       │   expansion=4.0,
        │   │   │       │   pre_norm=CoAtNorm.Config(
        │   │   │       │       channels_in=128,
        │   │   │       │       channels_out=128,
        │   │   │       │       momentum=0.1,
        │   │   │       │       eps=1e-05,
        │   │   │       │       elementwise_affine=True,
        │   │   │       │       device=None,
        │   │   │       │       dtype=None,
        │   │   │       │       activation=priml.model.coatnet.silu_inplace
        │   │   │       │   ),
        │   │   │       │   conv1_1x1=Conv2d.Config(
        │   │   │       │       channels_in=128,
        │   │   │       │       channels_out=512,
        │   │   │       │       kernel_size=1,
        │   │   │       │       stride=1,
        │   │   │       │       padding=0,
        │   │   │       │       dilation=1,
        │   │   │       │       groups=1,
        │   │   │       │       bias=False,
        │   │   │       │       depth_index=(),
        │   │   │       │       device=None,
        │   │   │       │       dtype=None,
        │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │       init_bias=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   norm1=CoAtNorm.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       momentum=0.1,
        │   │   │       │       eps=1e-05,
        │   │   │       │       elementwise_affine=True,
        │   │   │       │       device=None,
        │   │   │       │       dtype=None,
        │   │   │       │       activation=priml.model.coatnet.silu_inplace
        │   │   │       │   ),
        │   │   │       │   conv2_kxk=Conv2d.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       kernel_size=3,
        │   │   │       │       stride=2,
        │   │   │       │       padding=1,
        │   │   │       │       dilation=1,
        │   │   │       │       groups=512,
        │   │   │       │       bias=False,
        │   │   │       │       depth_index=(),
        │   │   │       │       device=None,
        │   │   │       │       dtype=None,
        │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │       init_bias=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   norm2=CoAtNorm.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       momentum=0.1,
        │   │   │       │       eps=1e-05,
        │   │   │       │       elementwise_affine=True,
        │   │   │       │       device=None,
        │   │   │       │       dtype=None,
        │   │   │       │       activation=priml.model.coatnet.silu_inplace
        │   │   │       │   ),
        │   │   │       │   se=SqueezeExcitation.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       channels_hidden=128,
        │   │   │       │       fc1=Conv2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=128,
        │   │   │       │           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_
        │   │   │       │       ),
        │   │   │       │       fc2=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=512,
        │   │   │       │           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_
        │   │   │       │       ),
        │   │   │       │       squeeze=functools.partial(<built-in method mean of type object at 0xdefacedeface>, dim=(2, 3), keepdim=True),
        │   │   │       │       expand=None,
        │   │   │       │       act=priml.model.coatnet.silu_inplace,
        │   │   │       │       gate=priml.model.swiglu.sigmoid
        │   │   │       │   ),
        │   │   │       │   conv3_1x1=Conv2d.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=256,
        │   │   │       │       kernel_size=1,
        │   │   │       │       stride=1,
        │   │   │       │       padding=0,
        │   │   │       │       dilation=1,
        │   │   │       │       groups=1,
        │   │   │       │       bias=False,
        │   │   │       │       depth_index=(),
        │   │   │       │       device=None,
        │   │   │       │       dtype=None,
        │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │       init_bias=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   shortcut=CoAtPool.Config(
        │   │   │       │       channels_in=128,
        │   │   │       │       channels_out=256,
        │   │   │       │       pool=AvgPool2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=2,
        │   │   │       │           stride=2,
        │   │   │       │           padding=0,
        │   │   │       │           ceil_mode=False,
        │   │   │       │           count_include_pad=True,
        │   │   │       │           divisor_override=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       )
        │   │   │       │   )
        │   │   │       ),
        │   │   │       block=[
        │   │   │       │   CoAtMBConv.Config(
        │   │   │       │       channels_in=256,
        │   │   │       │       channels_out=256,
        │   │   │       │       stride=1,
        │   │   │       │       expansion=4.0,
        │   │   │       │       pre_norm=CoAtNorm.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       conv1_1x1=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=CoAtNorm.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       conv2_kxk=Conv2d.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1_024,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=CoAtNorm.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       se=SqueezeExcitation.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           channels_hidden=256,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=1_024,
        │   │   │       │               channels_out=256,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=256,
        │   │   │       │               channels_out=1_024,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           squeeze=functools.partial(<built-in method mean of type object at 0xdefacedeface>, dim=(2, 3), keepdim=True),
        │   │   │       │           expand=None,
        │   │   │       │           act=priml.model.coatnet.silu_inplace,
        │   │   │       │           gate=priml.model.swiglu.sigmoid
        │   │   │       │       ),
        │   │   │       │       conv3_1x1=Conv2d.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       shortcut=CoAtPool.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           pool=AvgPool2d.Config(
        │   │   │       │               channels_in=256,
        │   │   │       │               channels_out=256,
        │   │   │       │               kernel_size=2,
        │   │   │       │               stride=2,
        │   │   │       │               padding=0,
        │   │   │       │               ceil_mode=False,
        │   │   │       │               count_include_pad=True,
        │   │   │       │               divisor_override=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=priml.model.init.kaiming_uniform,
        │   │   │       │               init_bias=torch.nn.init.zeros_
        │   │   │       │           )
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtMBConv.Config(
        │   │   │       │       channels_in=256,
        │   │   │       │       channels_out=256,
        │   │   │       │       stride=1,
        │   │   │       │       expansion=4.0,
        │   │   │       │       pre_norm=CoAtNorm.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       conv1_1x1=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=CoAtNorm.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       conv2_kxk=Conv2d.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1_024,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=CoAtNorm.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       se=SqueezeExcitation.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           channels_hidden=256,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=1_024,
        │   │   │       │               channels_out=256,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=256,
        │   │   │       │               channels_out=1_024,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           squeeze=functools.partial(<built-in method mean of type object at 0xdefacedeface>, dim=(2, 3), keepdim=True),
        │   │   │       │           expand=None,
        │   │   │       │           act=priml.model.coatnet.silu_inplace,
        │   │   │       │           gate=priml.model.swiglu.sigmoid
        │   │   │       │       ),
        │   │   │       │       conv3_1x1=Conv2d.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       shortcut=CoAtPool.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           pool=AvgPool2d.Config(
        │   │   │       │               channels_in=256,
        │   │   │       │               channels_out=256,
        │   │   │       │               kernel_size=2,
        │   │   │       │               stride=2,
        │   │   │       │               padding=0,
        │   │   │       │               ceil_mode=False,
        │   │   │       │               count_include_pad=True,
        │   │   │       │               divisor_override=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=priml.model.init.kaiming_uniform,
        │   │   │       │               init_bias=torch.nn.init.zeros_
        │   │   │       │           )
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtMBConv.Config(
        │   │   │       │       channels_in=256,
        │   │   │       │       channels_out=256,
        │   │   │       │       stride=1,
        │   │   │       │       expansion=4.0,
        │   │   │       │       pre_norm=CoAtNorm.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       conv1_1x1=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=CoAtNorm.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       conv2_kxk=Conv2d.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1_024,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=CoAtNorm.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       se=SqueezeExcitation.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           channels_hidden=256,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=1_024,
        │   │   │       │               channels_out=256,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=256,
        │   │   │       │               channels_out=1_024,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           squeeze=functools.partial(<built-in method mean of type object at 0xdefacedeface>, dim=(2, 3), keepdim=True),
        │   │   │       │           expand=None,
        │   │   │       │           act=priml.model.coatnet.silu_inplace,
        │   │   │       │           gate=priml.model.swiglu.sigmoid
        │   │   │       │       ),
        │   │   │       │       conv3_1x1=Conv2d.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       shortcut=CoAtPool.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           pool=AvgPool2d.Config(
        │   │   │       │               channels_in=256,
        │   │   │       │               channels_out=256,
        │   │   │       │               kernel_size=2,
        │   │   │       │               stride=2,
        │   │   │       │               padding=0,
        │   │   │       │               ceil_mode=False,
        │   │   │       │               count_include_pad=True,
        │   │   │       │               divisor_override=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=priml.model.init.kaiming_uniform,
        │   │   │       │               init_bias=torch.nn.init.zeros_
        │   │   │       │           )
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtMBConv.Config(
        │   │   │       │       channels_in=256,
        │   │   │       │       channels_out=256,
        │   │   │       │       stride=1,
        │   │   │       │       expansion=4.0,
        │   │   │       │       pre_norm=CoAtNorm.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       conv1_1x1=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=CoAtNorm.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       conv2_kxk=Conv2d.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1_024,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=CoAtNorm.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       se=SqueezeExcitation.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           channels_hidden=256,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=1_024,
        │   │   │       │               channels_out=256,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=256,
        │   │   │       │               channels_out=1_024,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           squeeze=functools.partial(<built-in method mean of type object at 0xdefacedeface>, dim=(2, 3), keepdim=True),
        │   │   │       │           expand=None,
        │   │   │       │           act=priml.model.coatnet.silu_inplace,
        │   │   │       │           gate=priml.model.swiglu.sigmoid
        │   │   │       │       ),
        │   │   │       │       conv3_1x1=Conv2d.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       shortcut=CoAtPool.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           pool=AvgPool2d.Config(
        │   │   │       │               channels_in=256,
        │   │   │       │               channels_out=256,
        │   │   │       │               kernel_size=2,
        │   │   │       │               stride=2,
        │   │   │       │               padding=0,
        │   │   │       │               ceil_mode=False,
        │   │   │       │               count_include_pad=True,
        │   │   │       │               divisor_override=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=priml.model.init.kaiming_uniform,
        │   │   │       │               init_bias=torch.nn.init.zeros_
        │   │   │       │           )
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtMBConv.Config(
        │   │   │       │       channels_in=256,
        │   │   │       │       channels_out=256,
        │   │   │       │       stride=1,
        │   │   │       │       expansion=4.0,
        │   │   │       │       pre_norm=CoAtNorm.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       conv1_1x1=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=CoAtNorm.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       conv2_kxk=Conv2d.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1_024,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=CoAtNorm.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           activation=priml.model.coatnet.silu_inplace
        │   │   │       │       ),
        │   │   │       │       se=SqueezeExcitation.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           channels_hidden=256,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=1_024,
        │   │   │       │               channels_out=256,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=256,
        │   │   │       │               channels_out=1_024,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           squeeze=functools.partial(<built-in method mean of type object at 0xdefacedeface>, dim=(2, 3), keepdim=True),
        │   │   │       │           expand=None,
        │   │   │       │           act=priml.model.coatnet.silu_inplace,
        │   │   │       │           gate=priml.model.swiglu.sigmoid
        │   │   │       │       ),
        │   │   │       │       conv3_1x1=Conv2d.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       shortcut=CoAtPool.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           pool=AvgPool2d.Config(
        │   │   │       │               channels_in=256,
        │   │   │       │               channels_out=256,
        │   │   │       │               kernel_size=2,
        │   │   │       │               stride=2,
        │   │   │       │               padding=0,
        │   │   │       │               ceil_mode=False,
        │   │   │       │               count_include_pad=True,
        │   │   │       │               divisor_override=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=priml.model.init.kaiming_uniform,
        │   │   │       │               init_bias=torch.nn.init.zeros_
        │   │   │       │           )
        │   │   │       │       )
        │   │   │       │   )
        │   │   │       ]
        │   │   │   ),
        │   │   │   CoAtStage.Config(
        │   │   │       channels_in=256,
        │   │   │       channels_out=512,
        │   │   │       feature_size=(14, 14),
        │   │   │       blocks_per_stage=14,
        │   │   │       transition=CoAtTransformer.Config(
        │   │   │       │   channels_in=256,
        │   │   │       │   channels_out=512,
        │   │   │       │   norm1=CoAtDownNorm.Config(
        │   │   │       │       channels_in=256,
        │   │   │       │       norm=LayerNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       down=CoAtPool.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           pool=AvgPool2d.Config(
        │   │   │       │               channels_in=256,
        │   │   │       │               channels_out=256,
        │   │   │       │               kernel_size=2,
        │   │   │       │               stride=2,
        │   │   │       │               padding=0,
        │   │   │       │               ceil_mode=False,
        │   │   │       │               count_include_pad=True,
        │   │   │       │               divisor_override=None
        │   │   │       │           ),
        │   │   │       │           projection=Conv2d.Config(
        │   │   │       │               channels_in=256,
        │   │   │       │               channels_out=256,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   shortcut=CoAtPool.Config(
        │   │   │       │       channels_in=256,
        │   │   │       │       channels_out=512,
        │   │   │       │       pool=AvgPool2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=2,
        │   │   │       │           stride=2,
        │   │   │       │           padding=0,
        │   │   │       │           ceil_mode=False,
        │   │   │       │           count_include_pad=True,
        │   │   │       │           divisor_override=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=512,
        │   │   │       │           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_
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   attention=CoAtAttention.Config(
        │   │   │       │   │   channels_in=256,
        │   │   │       │   │   channels_out=512,
        │   │   │       │   │   channels_head=32,
        │   │   │       │   │   qkv=Conv2d.Config(
        │   │   │       │   │       channels_in=256,
        │   │   │       │   │       channels_out=768,
        │   │   │       │   │       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_
        │   │   │       │   │   ),
        │   │   │       │   │   relative=CoAtRelativeBias.Config(
        │   │   │       │   │       feature_size=(14, 14),
        │   │   │       │   │       num_heads=8,
        │   │   │       │   │       mlp=CoAtLinearMLP.Config(
        │   │   │       │   │           channels_hidden=512,
        │   │   │       │   │           channels_out=8,
        │   │   │       │   │           fc1=Linear.Config(
        │   │   │       │   │               channels_in=2,
        │   │   │       │   │               channels_out=512,
        │   │   │       │   │               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
        │   │   │       │   │           ),
        │   │   │       │   │           fc2=Linear.Config(
        │   │   │       │   │               channels_in=512,
        │   │   │       │   │               channels_out=8,
        │   │   │       │   │               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
        │   │   │       │   │           ),
        │   │   │       │   │           dropout=Dropout.Config(
        │   │   │       │   │               channels_in=512,
        │   │   │       │   │               channels_out=512,
        │   │   │       │   │               p=0.125,
        │   │   │       │   │               inplace=False
        │   │   │       │   │           )
        │   │   │       │   │       )
        │   │   │       │   │   ),
        │   │   │       │   │   projection=Conv2d.Config(
        │   │   │       │   │       channels_in=256,
        │   │   │       │   │       channels_out=512,
        │   │   │       │   │       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_
        │   │   │       │   │   )
        │   │   │       │   ),
        │   │   │       │   ls1=CoAtScale.Config(channels_in=512, initial=1e-06),
        │   │   │       │   norm2=LayerNorm2d.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       eps=1e-06,
        │   │   │       │       elementwise_affine=True,
        │   │   │       │       device=None,
        │   │   │       │       dtype=None
        │   │   │       │   ),
        │   │   │       │   mlp=CoAtConvMLP.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       expansion=4.0,
        │   │   │       │       fc1=Conv2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=2_048,
        │   │   │       │           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_
        │   │   │       │       ),
        │   │   │       │       fc2=Conv2d.Config(
        │   │   │       │           channels_in=2_048,
        │   │   │       │           channels_out=512,
        │   │   │       │           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_
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   ls2=CoAtScale.Config(channels_in=512, initial=1e-06)
        │   │   │       ),
        │   │   │       block=[
        │   │   │       │   CoAtTransformer.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       norm1=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       shortcut=Identity.Config(
        │   │   │       │           channels_in=-1,
        │   │   │       │           channels_out=-1
        │   │   │       │       ),
        │   │   │       │       attention=CoAtAttention.Config(
        │   │   │       │       │   channels_in=512,
        │   │   │       │       │   channels_out=512,
        │   │   │       │       │   channels_head=32,
        │   │   │       │       │   qkv=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=1_536,
        │   │   │       │       │       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_
        │   │   │       │       │   ),
        │   │   │       │       │   relative=CoAtRelativeBias.Config(
        │   │   │       │       │       feature_size=(14, 14),
        │   │   │       │       │       num_heads=16,
        │   │   │       │       │       mlp=CoAtLinearMLP.Config(
        │   │   │       │       │           channels_hidden=512,
        │   │   │       │       │           channels_out=16,
        │   │   │       │       │           fc1=Linear.Config(
        │   │   │       │       │               channels_in=2,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           fc2=Linear.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=16,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           dropout=Dropout.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               p=0.125,
        │   │   │       │       │               inplace=False
        │   │   │       │       │           )
        │   │   │       │       │       )
        │   │   │       │       │   ),
        │   │   │       │       │   projection=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=512,
        │   │   │       │       │       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_
        │   │   │       │       │   )
        │   │   │       │       ),
        │   │   │       │       ls1=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       ),
        │   │   │       │       norm2=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       mlp=CoAtConvMLP.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           expansion=4.0,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=2_048,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=2_048,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       ),
        │   │   │       │       ls2=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtTransformer.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       norm1=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       shortcut=Identity.Config(
        │   │   │       │           channels_in=-1,
        │   │   │       │           channels_out=-1
        │   │   │       │       ),
        │   │   │       │       attention=CoAtAttention.Config(
        │   │   │       │       │   channels_in=512,
        │   │   │       │       │   channels_out=512,
        │   │   │       │       │   channels_head=32,
        │   │   │       │       │   qkv=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=1_536,
        │   │   │       │       │       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_
        │   │   │       │       │   ),
        │   │   │       │       │   relative=CoAtRelativeBias.Config(
        │   │   │       │       │       feature_size=(14, 14),
        │   │   │       │       │       num_heads=16,
        │   │   │       │       │       mlp=CoAtLinearMLP.Config(
        │   │   │       │       │           channels_hidden=512,
        │   │   │       │       │           channels_out=16,
        │   │   │       │       │           fc1=Linear.Config(
        │   │   │       │       │               channels_in=2,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           fc2=Linear.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=16,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           dropout=Dropout.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               p=0.125,
        │   │   │       │       │               inplace=False
        │   │   │       │       │           )
        │   │   │       │       │       )
        │   │   │       │       │   ),
        │   │   │       │       │   projection=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=512,
        │   │   │       │       │       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_
        │   │   │       │       │   )
        │   │   │       │       ),
        │   │   │       │       ls1=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       ),
        │   │   │       │       norm2=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       mlp=CoAtConvMLP.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           expansion=4.0,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=2_048,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=2_048,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       ),
        │   │   │       │       ls2=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtTransformer.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       norm1=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       shortcut=Identity.Config(
        │   │   │       │           channels_in=-1,
        │   │   │       │           channels_out=-1
        │   │   │       │       ),
        │   │   │       │       attention=CoAtAttention.Config(
        │   │   │       │       │   channels_in=512,
        │   │   │       │       │   channels_out=512,
        │   │   │       │       │   channels_head=32,
        │   │   │       │       │   qkv=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=1_536,
        │   │   │       │       │       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_
        │   │   │       │       │   ),
        │   │   │       │       │   relative=CoAtRelativeBias.Config(
        │   │   │       │       │       feature_size=(14, 14),
        │   │   │       │       │       num_heads=16,
        │   │   │       │       │       mlp=CoAtLinearMLP.Config(
        │   │   │       │       │           channels_hidden=512,
        │   │   │       │       │           channels_out=16,
        │   │   │       │       │           fc1=Linear.Config(
        │   │   │       │       │               channels_in=2,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           fc2=Linear.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=16,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           dropout=Dropout.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               p=0.125,
        │   │   │       │       │               inplace=False
        │   │   │       │       │           )
        │   │   │       │       │       )
        │   │   │       │       │   ),
        │   │   │       │       │   projection=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=512,
        │   │   │       │       │       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_
        │   │   │       │       │   )
        │   │   │       │       ),
        │   │   │       │       ls1=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       ),
        │   │   │       │       norm2=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       mlp=CoAtConvMLP.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           expansion=4.0,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=2_048,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=2_048,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       ),
        │   │   │       │       ls2=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtTransformer.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       norm1=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       shortcut=Identity.Config(
        │   │   │       │           channels_in=-1,
        │   │   │       │           channels_out=-1
        │   │   │       │       ),
        │   │   │       │       attention=CoAtAttention.Config(
        │   │   │       │       │   channels_in=512,
        │   │   │       │       │   channels_out=512,
        │   │   │       │       │   channels_head=32,
        │   │   │       │       │   qkv=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=1_536,
        │   │   │       │       │       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_
        │   │   │       │       │   ),
        │   │   │       │       │   relative=CoAtRelativeBias.Config(
        │   │   │       │       │       feature_size=(14, 14),
        │   │   │       │       │       num_heads=16,
        │   │   │       │       │       mlp=CoAtLinearMLP.Config(
        │   │   │       │       │           channels_hidden=512,
        │   │   │       │       │           channels_out=16,
        │   │   │       │       │           fc1=Linear.Config(
        │   │   │       │       │               channels_in=2,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           fc2=Linear.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=16,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           dropout=Dropout.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               p=0.125,
        │   │   │       │       │               inplace=False
        │   │   │       │       │           )
        │   │   │       │       │       )
        │   │   │       │       │   ),
        │   │   │       │       │   projection=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=512,
        │   │   │       │       │       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_
        │   │   │       │       │   )
        │   │   │       │       ),
        │   │   │       │       ls1=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       ),
        │   │   │       │       norm2=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       mlp=CoAtConvMLP.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           expansion=4.0,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=2_048,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=2_048,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       ),
        │   │   │       │       ls2=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtTransformer.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       norm1=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       shortcut=Identity.Config(
        │   │   │       │           channels_in=-1,
        │   │   │       │           channels_out=-1
        │   │   │       │       ),
        │   │   │       │       attention=CoAtAttention.Config(
        │   │   │       │       │   channels_in=512,
        │   │   │       │       │   channels_out=512,
        │   │   │       │       │   channels_head=32,
        │   │   │       │       │   qkv=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=1_536,
        │   │   │       │       │       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_
        │   │   │       │       │   ),
        │   │   │       │       │   relative=CoAtRelativeBias.Config(
        │   │   │       │       │       feature_size=(14, 14),
        │   │   │       │       │       num_heads=16,
        │   │   │       │       │       mlp=CoAtLinearMLP.Config(
        │   │   │       │       │           channels_hidden=512,
        │   │   │       │       │           channels_out=16,
        │   │   │       │       │           fc1=Linear.Config(
        │   │   │       │       │               channels_in=2,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           fc2=Linear.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=16,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           dropout=Dropout.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               p=0.125,
        │   │   │       │       │               inplace=False
        │   │   │       │       │           )
        │   │   │       │       │       )
        │   │   │       │       │   ),
        │   │   │       │       │   projection=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=512,
        │   │   │       │       │       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_
        │   │   │       │       │   )
        │   │   │       │       ),
        │   │   │       │       ls1=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       ),
        │   │   │       │       norm2=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       mlp=CoAtConvMLP.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           expansion=4.0,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=2_048,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=2_048,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       ),
        │   │   │       │       ls2=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtTransformer.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       norm1=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       shortcut=Identity.Config(
        │   │   │       │           channels_in=-1,
        │   │   │       │           channels_out=-1
        │   │   │       │       ),
        │   │   │       │       attention=CoAtAttention.Config(
        │   │   │       │       │   channels_in=512,
        │   │   │       │       │   channels_out=512,
        │   │   │       │       │   channels_head=32,
        │   │   │       │       │   qkv=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=1_536,
        │   │   │       │       │       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_
        │   │   │       │       │   ),
        │   │   │       │       │   relative=CoAtRelativeBias.Config(
        │   │   │       │       │       feature_size=(14, 14),
        │   │   │       │       │       num_heads=16,
        │   │   │       │       │       mlp=CoAtLinearMLP.Config(
        │   │   │       │       │           channels_hidden=512,
        │   │   │       │       │           channels_out=16,
        │   │   │       │       │           fc1=Linear.Config(
        │   │   │       │       │               channels_in=2,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           fc2=Linear.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=16,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           dropout=Dropout.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               p=0.125,
        │   │   │       │       │               inplace=False
        │   │   │       │       │           )
        │   │   │       │       │       )
        │   │   │       │       │   ),
        │   │   │       │       │   projection=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=512,
        │   │   │       │       │       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_
        │   │   │       │       │   )
        │   │   │       │       ),
        │   │   │       │       ls1=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       ),
        │   │   │       │       norm2=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       mlp=CoAtConvMLP.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           expansion=4.0,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=2_048,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=2_048,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       ),
        │   │   │       │       ls2=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtTransformer.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       norm1=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       shortcut=Identity.Config(
        │   │   │       │           channels_in=-1,
        │   │   │       │           channels_out=-1
        │   │   │       │       ),
        │   │   │       │       attention=CoAtAttention.Config(
        │   │   │       │       │   channels_in=512,
        │   │   │       │       │   channels_out=512,
        │   │   │       │       │   channels_head=32,
        │   │   │       │       │   qkv=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=1_536,
        │   │   │       │       │       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_
        │   │   │       │       │   ),
        │   │   │       │       │   relative=CoAtRelativeBias.Config(
        │   │   │       │       │       feature_size=(14, 14),
        │   │   │       │       │       num_heads=16,
        │   │   │       │       │       mlp=CoAtLinearMLP.Config(
        │   │   │       │       │           channels_hidden=512,
        │   │   │       │       │           channels_out=16,
        │   │   │       │       │           fc1=Linear.Config(
        │   │   │       │       │               channels_in=2,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           fc2=Linear.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=16,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           dropout=Dropout.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               p=0.125,
        │   │   │       │       │               inplace=False
        │   │   │       │       │           )
        │   │   │       │       │       )
        │   │   │       │       │   ),
        │   │   │       │       │   projection=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=512,
        │   │   │       │       │       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_
        │   │   │       │       │   )
        │   │   │       │       ),
        │   │   │       │       ls1=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       ),
        │   │   │       │       norm2=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       mlp=CoAtConvMLP.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           expansion=4.0,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=2_048,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=2_048,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       ),
        │   │   │       │       ls2=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtTransformer.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       norm1=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       shortcut=Identity.Config(
        │   │   │       │           channels_in=-1,
        │   │   │       │           channels_out=-1
        │   │   │       │       ),
        │   │   │       │       attention=CoAtAttention.Config(
        │   │   │       │       │   channels_in=512,
        │   │   │       │       │   channels_out=512,
        │   │   │       │       │   channels_head=32,
        │   │   │       │       │   qkv=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=1_536,
        │   │   │       │       │       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_
        │   │   │       │       │   ),
        │   │   │       │       │   relative=CoAtRelativeBias.Config(
        │   │   │       │       │       feature_size=(14, 14),
        │   │   │       │       │       num_heads=16,
        │   │   │       │       │       mlp=CoAtLinearMLP.Config(
        │   │   │       │       │           channels_hidden=512,
        │   │   │       │       │           channels_out=16,
        │   │   │       │       │           fc1=Linear.Config(
        │   │   │       │       │               channels_in=2,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           fc2=Linear.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=16,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           dropout=Dropout.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               p=0.125,
        │   │   │       │       │               inplace=False
        │   │   │       │       │           )
        │   │   │       │       │       )
        │   │   │       │       │   ),
        │   │   │       │       │   projection=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=512,
        │   │   │       │       │       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_
        │   │   │       │       │   )
        │   │   │       │       ),
        │   │   │       │       ls1=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       ),
        │   │   │       │       norm2=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       mlp=CoAtConvMLP.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           expansion=4.0,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=2_048,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=2_048,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       ),
        │   │   │       │       ls2=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtTransformer.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       norm1=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       shortcut=Identity.Config(
        │   │   │       │           channels_in=-1,
        │   │   │       │           channels_out=-1
        │   │   │       │       ),
        │   │   │       │       attention=CoAtAttention.Config(
        │   │   │       │       │   channels_in=512,
        │   │   │       │       │   channels_out=512,
        │   │   │       │       │   channels_head=32,
        │   │   │       │       │   qkv=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=1_536,
        │   │   │       │       │       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_
        │   │   │       │       │   ),
        │   │   │       │       │   relative=CoAtRelativeBias.Config(
        │   │   │       │       │       feature_size=(14, 14),
        │   │   │       │       │       num_heads=16,
        │   │   │       │       │       mlp=CoAtLinearMLP.Config(
        │   │   │       │       │           channels_hidden=512,
        │   │   │       │       │           channels_out=16,
        │   │   │       │       │           fc1=Linear.Config(
        │   │   │       │       │               channels_in=2,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           fc2=Linear.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=16,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           dropout=Dropout.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               p=0.125,
        │   │   │       │       │               inplace=False
        │   │   │       │       │           )
        │   │   │       │       │       )
        │   │   │       │       │   ),
        │   │   │       │       │   projection=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=512,
        │   │   │       │       │       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_
        │   │   │       │       │   )
        │   │   │       │       ),
        │   │   │       │       ls1=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       ),
        │   │   │       │       norm2=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       mlp=CoAtConvMLP.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           expansion=4.0,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=2_048,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=2_048,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       ),
        │   │   │       │       ls2=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtTransformer.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       norm1=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       shortcut=Identity.Config(
        │   │   │       │           channels_in=-1,
        │   │   │       │           channels_out=-1
        │   │   │       │       ),
        │   │   │       │       attention=CoAtAttention.Config(
        │   │   │       │       │   channels_in=512,
        │   │   │       │       │   channels_out=512,
        │   │   │       │       │   channels_head=32,
        │   │   │       │       │   qkv=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=1_536,
        │   │   │       │       │       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_
        │   │   │       │       │   ),
        │   │   │       │       │   relative=CoAtRelativeBias.Config(
        │   │   │       │       │       feature_size=(14, 14),
        │   │   │       │       │       num_heads=16,
        │   │   │       │       │       mlp=CoAtLinearMLP.Config(
        │   │   │       │       │           channels_hidden=512,
        │   │   │       │       │           channels_out=16,
        │   │   │       │       │           fc1=Linear.Config(
        │   │   │       │       │               channels_in=2,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           fc2=Linear.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=16,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           dropout=Dropout.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               p=0.125,
        │   │   │       │       │               inplace=False
        │   │   │       │       │           )
        │   │   │       │       │       )
        │   │   │       │       │   ),
        │   │   │       │       │   projection=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=512,
        │   │   │       │       │       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_
        │   │   │       │       │   )
        │   │   │       │       ),
        │   │   │       │       ls1=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       ),
        │   │   │       │       norm2=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       mlp=CoAtConvMLP.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           expansion=4.0,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=2_048,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=2_048,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       ),
        │   │   │       │       ls2=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtTransformer.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       norm1=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       shortcut=Identity.Config(
        │   │   │       │           channels_in=-1,
        │   │   │       │           channels_out=-1
        │   │   │       │       ),
        │   │   │       │       attention=CoAtAttention.Config(
        │   │   │       │       │   channels_in=512,
        │   │   │       │       │   channels_out=512,
        │   │   │       │       │   channels_head=32,
        │   │   │       │       │   qkv=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=1_536,
        │   │   │       │       │       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_
        │   │   │       │       │   ),
        │   │   │       │       │   relative=CoAtRelativeBias.Config(
        │   │   │       │       │       feature_size=(14, 14),
        │   │   │       │       │       num_heads=16,
        │   │   │       │       │       mlp=CoAtLinearMLP.Config(
        │   │   │       │       │           channels_hidden=512,
        │   │   │       │       │           channels_out=16,
        │   │   │       │       │           fc1=Linear.Config(
        │   │   │       │       │               channels_in=2,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           fc2=Linear.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=16,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           dropout=Dropout.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               p=0.125,
        │   │   │       │       │               inplace=False
        │   │   │       │       │           )
        │   │   │       │       │       )
        │   │   │       │       │   ),
        │   │   │       │       │   projection=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=512,
        │   │   │       │       │       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_
        │   │   │       │       │   )
        │   │   │       │       ),
        │   │   │       │       ls1=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       ),
        │   │   │       │       norm2=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       mlp=CoAtConvMLP.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           expansion=4.0,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=2_048,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=2_048,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       ),
        │   │   │       │       ls2=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtTransformer.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       norm1=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       shortcut=Identity.Config(
        │   │   │       │           channels_in=-1,
        │   │   │       │           channels_out=-1
        │   │   │       │       ),
        │   │   │       │       attention=CoAtAttention.Config(
        │   │   │       │       │   channels_in=512,
        │   │   │       │       │   channels_out=512,
        │   │   │       │       │   channels_head=32,
        │   │   │       │       │   qkv=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=1_536,
        │   │   │       │       │       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_
        │   │   │       │       │   ),
        │   │   │       │       │   relative=CoAtRelativeBias.Config(
        │   │   │       │       │       feature_size=(14, 14),
        │   │   │       │       │       num_heads=16,
        │   │   │       │       │       mlp=CoAtLinearMLP.Config(
        │   │   │       │       │           channels_hidden=512,
        │   │   │       │       │           channels_out=16,
        │   │   │       │       │           fc1=Linear.Config(
        │   │   │       │       │               channels_in=2,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           fc2=Linear.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=16,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           dropout=Dropout.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               p=0.125,
        │   │   │       │       │               inplace=False
        │   │   │       │       │           )
        │   │   │       │       │       )
        │   │   │       │       │   ),
        │   │   │       │       │   projection=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=512,
        │   │   │       │       │       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_
        │   │   │       │       │   )
        │   │   │       │       ),
        │   │   │       │       ls1=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       ),
        │   │   │       │       norm2=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       mlp=CoAtConvMLP.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           expansion=4.0,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=2_048,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=2_048,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       ),
        │   │   │       │       ls2=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   CoAtTransformer.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       norm1=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       shortcut=Identity.Config(
        │   │   │       │           channels_in=-1,
        │   │   │       │           channels_out=-1
        │   │   │       │       ),
        │   │   │       │       attention=CoAtAttention.Config(
        │   │   │       │       │   channels_in=512,
        │   │   │       │       │   channels_out=512,
        │   │   │       │       │   channels_head=32,
        │   │   │       │       │   qkv=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=1_536,
        │   │   │       │       │       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_
        │   │   │       │       │   ),
        │   │   │       │       │   relative=CoAtRelativeBias.Config(
        │   │   │       │       │       feature_size=(14, 14),
        │   │   │       │       │       num_heads=16,
        │   │   │       │       │       mlp=CoAtLinearMLP.Config(
        │   │   │       │       │           channels_hidden=512,
        │   │   │       │       │           channels_out=16,
        │   │   │       │       │           fc1=Linear.Config(
        │   │   │       │       │               channels_in=2,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           fc2=Linear.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=16,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           dropout=Dropout.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               p=0.125,
        │   │   │       │       │               inplace=False
        │   │   │       │       │           )
        │   │   │       │       │       )
        │   │   │       │       │   ),
        │   │   │       │       │   projection=Conv2d.Config(
        │   │   │       │       │       channels_in=512,
        │   │   │       │       │       channels_out=512,
        │   │   │       │       │       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_
        │   │   │       │       │   )
        │   │   │       │       ),
        │   │   │       │       ls1=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       ),
        │   │   │       │       norm2=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       mlp=CoAtConvMLP.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           expansion=4.0,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=2_048,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=2_048,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       ),
        │   │   │       │       ls2=CoAtScale.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           initial=1e-06
        │   │   │       │       )
        │   │   │       │   )
        │   │   │       ]
        │   │   │   ),
        │   │   │   CoAtStage.Config(
        │   │   │       channels_in=512,
        │   │   │       channels_out=1_024,
        │   │   │       feature_size=(7, 7),
        │   │   │       blocks_per_stage=2,
        │   │   │       transition=CoAtTransformer.Config(
        │   │   │       │   channels_in=512,
        │   │   │       │   channels_out=1_024,
        │   │   │       │   norm1=CoAtDownNorm.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       norm=LayerNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       down=CoAtPool.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           pool=AvgPool2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=512,
        │   │   │       │               kernel_size=2,
        │   │   │       │               stride=2,
        │   │   │       │               padding=0,
        │   │   │       │               ceil_mode=False,
        │   │   │       │               count_include_pad=True,
        │   │   │       │               divisor_override=None
        │   │   │       │           ),
        │   │   │       │           projection=Conv2d.Config(
        │   │   │       │               channels_in=512,
        │   │   │       │               channels_out=512,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   shortcut=CoAtPool.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=1_024,
        │   │   │       │       pool=AvgPool2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           kernel_size=2,
        │   │   │       │           stride=2,
        │   │   │       │           padding=0,
        │   │   │       │           ceil_mode=False,
        │   │   │       │           count_include_pad=True,
        │   │   │       │           divisor_override=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           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_
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   attention=CoAtAttention.Config(
        │   │   │       │   │   channels_in=512,
        │   │   │       │   │   channels_out=1_024,
        │   │   │       │   │   channels_head=32,
        │   │   │       │   │   qkv=Conv2d.Config(
        │   │   │       │   │       channels_in=512,
        │   │   │       │   │       channels_out=1_536,
        │   │   │       │   │       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_
        │   │   │       │   │   ),
        │   │   │       │   │   relative=CoAtRelativeBias.Config(
        │   │   │       │   │       feature_size=(7, 7),
        │   │   │       │   │       num_heads=16,
        │   │   │       │   │       mlp=CoAtLinearMLP.Config(
        │   │   │       │   │           channels_hidden=512,
        │   │   │       │   │           channels_out=16,
        │   │   │       │   │           fc1=Linear.Config(
        │   │   │       │   │               channels_in=2,
        │   │   │       │   │               channels_out=512,
        │   │   │       │   │               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
        │   │   │       │   │           ),
        │   │   │       │   │           fc2=Linear.Config(
        │   │   │       │   │               channels_in=512,
        │   │   │       │   │               channels_out=16,
        │   │   │       │   │               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
        │   │   │       │   │           ),
        │   │   │       │   │           dropout=Dropout.Config(
        │   │   │       │   │               channels_in=512,
        │   │   │       │   │               channels_out=512,
        │   │   │       │   │               p=0.125,
        │   │   │       │   │               inplace=False
        │   │   │       │   │           )
        │   │   │       │   │       )
        │   │   │       │   │   ),
        │   │   │       │   │   projection=Conv2d.Config(
        │   │   │       │   │       channels_in=512,
        │   │   │       │   │       channels_out=1_024,
        │   │   │       │   │       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_
        │   │   │       │   │   )
        │   │   │       │   ),
        │   │   │       │   ls1=CoAtScale.Config(channels_in=1024, initial=1e-06),
        │   │   │       │   norm2=LayerNorm2d.Config(
        │   │   │       │       channels_in=1_024,
        │   │   │       │       channels_out=1_024,
        │   │   │       │       eps=1e-06,
        │   │   │       │       elementwise_affine=True,
        │   │   │       │       device=None,
        │   │   │       │       dtype=None
        │   │   │       │   ),
        │   │   │       │   mlp=CoAtConvMLP.Config(
        │   │   │       │       channels_in=1_024,
        │   │   │       │       expansion=4.0,
        │   │   │       │       fc1=Conv2d.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=4_096,
        │   │   │       │           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_
        │   │   │       │       ),
        │   │   │       │       fc2=Conv2d.Config(
        │   │   │       │           channels_in=4_096,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           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_
        │   │   │       │       )
        │   │   │       │   ),
        │   │   │       │   ls2=CoAtScale.Config(channels_in=1024, initial=1e-06)
        │   │   │       ),
        │   │   │       block=[
        │   │   │       │   CoAtTransformer.Config(
        │   │   │       │       channels_in=1_024,
        │   │   │       │       channels_out=1_024,
        │   │   │       │       norm1=LayerNorm2d.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       shortcut=Identity.Config(
        │   │   │       │           channels_in=-1,
        │   │   │       │           channels_out=-1
        │   │   │       │       ),
        │   │   │       │       attention=CoAtAttention.Config(
        │   │   │       │       │   channels_in=1_024,
        │   │   │       │       │   channels_out=1_024,
        │   │   │       │       │   channels_head=32,
        │   │   │       │       │   qkv=Conv2d.Config(
        │   │   │       │       │       channels_in=1_024,
        │   │   │       │       │       channels_out=3_072,
        │   │   │       │       │       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_
        │   │   │       │       │   ),
        │   │   │       │       │   relative=CoAtRelativeBias.Config(
        │   │   │       │       │       feature_size=(7, 7),
        │   │   │       │       │       num_heads=32,
        │   │   │       │       │       mlp=CoAtLinearMLP.Config(
        │   │   │       │       │           channels_hidden=512,
        │   │   │       │       │           channels_out=32,
        │   │   │       │       │           fc1=Linear.Config(
        │   │   │       │       │               channels_in=2,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           fc2=Linear.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=32,
        │   │   │       │       │               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
        │   │   │       │       │           ),
        │   │   │       │       │           dropout=Dropout.Config(
        │   │   │       │       │               channels_in=512,
        │   │   │       │       │               channels_out=512,
        │   │   │       │       │               p=0.125,
        │   │   │       │       │               inplace=False
        │   │   │       │       │           )
        │   │   │       │       │       )
        │   │   │       │       │   ),
        │   │   │       │       │   projection=Conv2d.Config(
        │   │   │       │       │       channels_in=1_024,
        │   │   │       │       │       channels_out=1_024,
        │   │   │       │       │       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_
        │   │   │       │       │   )
        │   │   │       │       ),
        │   │   │       │       ls1=CoAtScale.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           initial=1e-06
        │   │   │       │       ),
        │   │   │       │       norm2=LayerNorm2d.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           channels_out=1_024,
        │   │   │       │           eps=1e-06,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       mlp=CoAtConvMLP.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           expansion=4.0,
        │   │   │       │           fc1=Conv2d.Config(
        │   │   │       │               channels_in=1_024,
        │   │   │       │               channels_out=4_096,
        │   │   │       │               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_
        │   │   │       │           ),
        │   │   │       │           fc2=Conv2d.Config(
        │   │   │       │               channels_in=4_096,
        │   │   │       │               channels_out=1_024,
        │   │   │       │               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_
        │   │   │       │           )
        │   │   │       │       ),
        │   │   │       │       ls2=CoAtScale.Config(
        │   │   │       │           channels_in=1_024,
        │   │   │       │           initial=1e-06
        │   │   │       │       )
        │   │   │       │   )
        │   │   │       ]
        │   │   │   )
        │   │   ],
        │   │   norm=LayerNorm2d.Config(
        │   │       channels_in=1_024,
        │   │       channels_out=1_024,
        │   │       eps=1e-06,
        │   │       elementwise_affine=True,
        │   │       device=None,
        │   │       dtype=None
        │   │   )
        │   ),
        │   head=FindingAttention.Config(
        │   │   channels_in=1_024,
        │   │   channels_out=12,
        │   │   channels_hidden=256,
        │   │   norm=LayerNorm.Config(
        │   │       channels_in=1_024,
        │   │       channels_out=1_024,
        │   │       eps=1e-05,
        │   │       elementwise_affine=True,
        │   │       device=None,
        │   │       dtype=None
        │   │   ),
        │   │   proj_in=Linear.Config(
        │   │       channels_in=1_024,
        │   │       channels_out=256,
        │   │       bias=True,
        │   │       depth_index=(),
        │   │       device=None,
        │   │       dtype=None,
        │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │       init_bias=None,
        │   │       linear_fn=<built-in function linear>,
        │   │       shard=None
        │   │   ),
        │   │   activation=Activation.Config(
        │   │       channels_in=256,
        │   │       channels_out=256,
        │   │       act=priml.baselines.rsna_knee.model.tanh
        │   │   ),
        │   │   dropout=Dropout.Config(
        │   │       channels_in=256,
        │   │       channels_out=256,
        │   │       p=0.2,
        │   │       inplace=False
        │   │   ),
        │   │   proj_out=Linear.Config(
        │   │       channels_in=256,
        │   │       channels_out=12,
        │   │       bias=True,
        │   │       depth_index=(),
        │   │       device=None,
        │   │       dtype=None,
        │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │       init_bias=None,
        │   │       linear_fn=<built-in function linear>,
        │   │       shard=None
        │   │   ),
        │   │   init_weight=functools.partial(torch.nn.init.trunc_normal_, std=0.02),
        │   │   init_bias=torch.nn.init.zeros_
        │   ),
        │   chunk_size=inf
        ),
        preparation=RaptorPrepare.Config(
            image_size=384,
            crop_mm=140.0,
            slots=[
                ('Sagittal', 1, 26),
                ('Sagittal', 0, 22),
                ('Coronal', 1, 18),
                ('Coronal', 0, 12),
                ('Axial', -1, 18)
            ]
        ),
        windows=RaptorWindows.Config(
            windows=94,
            crop_size=307,
            mean=(0.485, 0.456, 0.406),
            std=(0.229, 0.224, 0.225)
        ),
        views=[priml.baselines.rsna_knee.imaging.unchanged],
        score=RaptorScore.Config(
            device='cuda',
            dtypes_autocast=[torch.float16, None],
            check_logits=priml.baselines.rsna_knee.inference.accept_logits
        ),
        time_limit=7_200,
        cudnn_benchmark=True,
        allow_tf32_cudnn=True,
        allow_tf32_matmul=True
    )
