PublicRecipe.Config(
        study_name='rsna-knee',
        experiment_name='exp000',
        base_dir='/opt/scratch',
        working_dir=PosixPath('/opt/scratch/runs/rsna-knee/exp000'),
        data=StudyData.Config(
            base_dir='/opt/scratch',
            input_dir=PosixPath('/opt/scratch/datasets/rsna-knee/raw'),
            split='test'
        ),
        raptor384=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_r384_swa.pt'),
        │           sha256='7e5315dad125b99fc65b340b3de41de628e9be51ff5835355dd61c86472244ef'
        │       ),
        │       key='model',
        │       weights_only=False
        │   ),
        │   model=Raptor.Config(
        │   │   backbone=CoAtNet.Config(
        │   │   │   image_size=(320, 320),
        │   │   │   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=(80, 80),
        │   │   │   │       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=(40, 40),
        │   │   │   │       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=(20, 20),
        │   │   │   │       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=(20, 20),
        │   │   │   │       │   │       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=(20, 20),
        │   │   │   │       │       │       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=(20, 20),
        │   │   │   │       │       │       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=(20, 20),
        │   │   │   │       │       │       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=(20, 20),
        │   │   │   │       │       │       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=(20, 20),
        │   │   │   │       │       │       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=(20, 20),
        │   │   │   │       │       │       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=(20, 20),
        │   │   │   │       │       │       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=(20, 20),
        │   │   │   │       │       │       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=(20, 20),
        │   │   │   │       │       │       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=(20, 20),
        │   │   │   │       │       │       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=(20, 20),
        │   │   │   │       │       │       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=(20, 20),
        │   │   │   │       │       │       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=(20, 20),
        │   │   │   │       │       │       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=(10, 10),
        │   │   │   │       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=(10, 10),
        │   │   │   │       │   │       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=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
        │   │   │   │       │   )
        │   │   │   │       ),
        │   │   │   │       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=(10, 10),
        │   │   │   │       │       │       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=16
        │   ),
        │   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=320,
        │       mean=(0.485, 0.456, 0.406),
        │       std=(0.229, 0.224, 0.225)
        │   ),
        │   views=[
        │       priml.baselines.rsna_knee.imaging.unchanged,
        │       priml.baselines.rsna_knee.imaging.anatomical_mirror
        │   ],
        │   score=RaptorScore.Config(
        │       device='cuda',
        │       dtypes_autocast=[torch.float16],
        │       check_logits=priml.baselines.rsna_knee.inference.require_finite
        │   ),
        │   time_limit=7_200,
        │   cudnn_benchmark=True,
        │   allow_tf32_cudnn=True,
        │   allow_tf32_matmul=True
        ),
        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=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
        │   │   │   │       │   )
        │   │   │   │       ),
        │   │   │   │       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
        ),
        r7=R7Stage.Config(
        │   base_dir='/opt/scratch',
        │   checkpoints=[
        │       ReleasedState.Config(
        │           base_dir='/opt/scratch',
        │           file=LocalFile.Config(
        │               path=PosixPath('/opt/scratch/datasets/rsna-knee-r7-v2/E_lrprior_s0_f0.pt'),
        │               sha256='ca8a76e024828793e0663e5a2a13461de17e820bed6789eafe3f155072d3612f'
        │           ),
        │           key=None,
        │           weights_only=True
        │       ),
        │       ReleasedState.Config(
        │           base_dir='/opt/scratch',
        │           file=LocalFile.Config(
        │               path=PosixPath('/opt/scratch/datasets/rsna-knee-r7-v2/E_lrprior_s0_f1.pt'),
        │               sha256='fb432cf544f218608dbaecf7a8dd3c8b6a26b34a3dfe2dc198172c3e222d5910'
        │           ),
        │           key=None,
        │           weights_only=True
        │       ),
        │       ReleasedState.Config(
        │           base_dir='/opt/scratch',
        │           file=LocalFile.Config(
        │               path=PosixPath('/opt/scratch/datasets/rsna-knee-r7-v2/E_lrprior_s0_f2.pt'),
        │               sha256='a9718ae111d6307492e724eb1767f0b5a77397d9a4d82c50c1e146b046523bf7'
        │           ),
        │           key=None,
        │           weights_only=True
        │       ),
        │       ReleasedState.Config(
        │           base_dir='/opt/scratch',
        │           file=LocalFile.Config(
        │               path=PosixPath('/opt/scratch/datasets/rsna-knee-r7-v2/E_lrprior_s0_f3.pt'),
        │               sha256='073c22c71a9946c1457e37e56fec145f073a921da106b55a10884da01a40fc94'
        │           ),
        │           key=None,
        │           weights_only=True
        │       )
        │   ],
        │   model=R7.Config(
        │   │   channels_hidden=384,
        │   │   encoder=ConvNeXt.Config(
        │   │   │   channels_in=3,
        │   │   │   channels_hidden=(96, 192, 384, 768),
        │   │   │   blocks_per_stage=(3, 3, 27, 3),
        │   │   │   stem=ConvNeXtStem.Config(
        │   │   │       channels_in=3,
        │   │   │       channels_out=96,
        │   │   │       conv=Conv2d.Config(
        │   │   │           channels_in=3,
        │   │   │           channels_out=96,
        │   │   │           kernel_size=4,
        │   │   │           stride=4,
        │   │   │           padding=0,
        │   │   │           dilation=1,
        │   │   │           groups=1,
        │   │   │           bias=True,
        │   │   │           depth_index=(),
        │   │   │           device=None,
        │   │   │           dtype=None,
        │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │           init_bias=torch.nn.init.zeros_
        │   │   │       ),
        │   │   │       norm=LayerNorm2d.Config(
        │   │   │           channels_in=96,
        │   │   │           channels_out=96,
        │   │   │           eps=1e-06,
        │   │   │           elementwise_affine=True,
        │   │   │           device=None,
        │   │   │           dtype=None
        │   │   │       )
        │   │   │   ),
        │   │   │   stage=[
        │   │   │   │   ConvNeXtStage.Config(
        │   │   │   │       channels_in=96,
        │   │   │   │       channels_out=96,
        │   │   │   │       blocks_per_stage=2,
        │   │   │   │       downsample=Identity.Config(channels_in=96, channels_out=96),
        │   │   │   │       block=[
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=96,
        │   │   │   │       │       channels_out=96,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=96,
        │   │   │   │       │           channels_out=96,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=96,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=96,
        │   │   │   │       │           channels_out=96,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=96,
        │   │   │   │       │           channels_out=96,
        │   │   │   │       │           channels_hidden=384,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.0
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=96,
        │   │   │   │       │       channels_out=96,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=96,
        │   │   │   │       │           channels_out=96,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=96,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=96,
        │   │   │   │       │           channels_out=96,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=96,
        │   │   │   │       │           channels_out=96,
        │   │   │   │       │           channels_hidden=384,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.0028571428265422583
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=96,
        │   │   │   │       │       channels_out=96,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=96,
        │   │   │   │       │           channels_out=96,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=96,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=96,
        │   │   │   │       │           channels_out=96,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=96,
        │   │   │   │       │           channels_out=96,
        │   │   │   │       │           channels_hidden=384,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.0057142856530845165
        │   │   │   │       │   )
        │   │   │   │       ]
        │   │   │   │   ),
        │   │   │   │   ConvNeXtStage.Config(
        │   │   │   │       channels_in=96,
        │   │   │   │       channels_out=192,
        │   │   │   │       blocks_per_stage=2,
        │   │   │   │       downsample=Sequential.Config(
        │   │   │   │           channels_in=96,
        │   │   │   │           channels_out=192,
        │   │   │   │           elements=[
        │   │   │   │               LayerNorm2d.Config(
        │   │   │   │                   channels_in=96,
        │   │   │   │                   channels_out=96,
        │   │   │   │                   eps=1e-06,
        │   │   │   │                   elementwise_affine=True,
        │   │   │   │                   device=None,
        │   │   │   │                   dtype=None
        │   │   │   │               ),
        │   │   │   │               Conv2d.Config(
        │   │   │   │                   channels_in=96,
        │   │   │   │                   channels_out=192,
        │   │   │   │                   kernel_size=2,
        │   │   │   │                   stride=2,
        │   │   │   │                   padding=0,
        │   │   │   │                   dilation=1,
        │   │   │   │                   groups=1,
        │   │   │   │                   bias=True,
        │   │   │   │                   depth_index=(),
        │   │   │   │                   device=None,
        │   │   │   │                   dtype=None,
        │   │   │   │                   init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │                   init_bias=torch.nn.init.zeros_
        │   │   │   │               )
        │   │   │   │           ],
        │   │   │   │           repeat=1,
        │   │   │   │           depth_index=()
        │   │   │   │       ),
        │   │   │   │       block=[
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=192,
        │   │   │   │       │       channels_out=192,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=192,
        │   │   │   │       │           channels_out=192,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=192,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=192,
        │   │   │   │       │           channels_out=192,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=192,
        │   │   │   │       │           channels_out=192,
        │   │   │   │       │           channels_hidden=768,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.008571428246796131
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=192,
        │   │   │   │       │       channels_out=192,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=192,
        │   │   │   │       │           channels_out=192,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=192,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=192,
        │   │   │   │       │           channels_out=192,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=192,
        │   │   │   │       │           channels_out=192,
        │   │   │   │       │           channels_hidden=768,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.011428571306169033
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=192,
        │   │   │   │       │       channels_out=192,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=192,
        │   │   │   │       │           channels_out=192,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=192,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=192,
        │   │   │   │       │           channels_out=192,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=192,
        │   │   │   │       │           channels_out=192,
        │   │   │   │       │           channels_hidden=768,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.014285714365541935
        │   │   │   │       │   )
        │   │   │   │       ]
        │   │   │   │   ),
        │   │   │   │   ConvNeXtStage.Config(
        │   │   │   │       channels_in=192,
        │   │   │   │       channels_out=384,
        │   │   │   │       blocks_per_stage=2,
        │   │   │   │       downsample=Sequential.Config(
        │   │   │   │           channels_in=192,
        │   │   │   │           channels_out=384,
        │   │   │   │           elements=[
        │   │   │   │               LayerNorm2d.Config(
        │   │   │   │                   channels_in=192,
        │   │   │   │                   channels_out=192,
        │   │   │   │                   eps=1e-06,
        │   │   │   │                   elementwise_affine=True,
        │   │   │   │                   device=None,
        │   │   │   │                   dtype=None
        │   │   │   │               ),
        │   │   │   │               Conv2d.Config(
        │   │   │   │                   channels_in=192,
        │   │   │   │                   channels_out=384,
        │   │   │   │                   kernel_size=2,
        │   │   │   │                   stride=2,
        │   │   │   │                   padding=0,
        │   │   │   │                   dilation=1,
        │   │   │   │                   groups=1,
        │   │   │   │                   bias=True,
        │   │   │   │                   depth_index=(),
        │   │   │   │                   device=None,
        │   │   │   │                   dtype=None,
        │   │   │   │                   init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │                   init_bias=torch.nn.init.zeros_
        │   │   │   │               )
        │   │   │   │           ],
        │   │   │   │           repeat=1,
        │   │   │   │           depth_index=()
        │   │   │   │       ),
        │   │   │   │       block=[
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.017142856493592262
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.019999999552965164
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.022857142612338066
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.025714285671710968
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.02857142873108387
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.03142857179045677
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.034285712987184525
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.037142857909202576
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.03999999910593033
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.04285714402794838
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.04571428522467613
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.048571426421403885
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.051428575068712234
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.05428571626543999
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.05714285746216774
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.06000000238418579
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.06285714358091354
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.0657142847776413
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.06857143342494965
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.0714285746216774
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.07428571581840515
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.0771428570151329
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.07999999821186066
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.082857146859169
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.08571428805589676
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.08857142925262451
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=384,
        │   │   │   │       │       channels_out=384,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=384,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=384,
        │   │   │   │       │           channels_out=384,
        │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.09142857044935226
        │   │   │   │       │   )
        │   │   │   │       ]
        │   │   │   │   ),
        │   │   │   │   ConvNeXtStage.Config(
        │   │   │   │       channels_in=384,
        │   │   │   │       channels_out=768,
        │   │   │   │       blocks_per_stage=2,
        │   │   │   │       downsample=Sequential.Config(
        │   │   │   │           channels_in=384,
        │   │   │   │           channels_out=768,
        │   │   │   │           elements=[
        │   │   │   │               LayerNorm2d.Config(
        │   │   │   │                   channels_in=384,
        │   │   │   │                   channels_out=384,
        │   │   │   │                   eps=1e-06,
        │   │   │   │                   elementwise_affine=True,
        │   │   │   │                   device=None,
        │   │   │   │                   dtype=None
        │   │   │   │               ),
        │   │   │   │               Conv2d.Config(
        │   │   │   │                   channels_in=384,
        │   │   │   │                   channels_out=768,
        │   │   │   │                   kernel_size=2,
        │   │   │   │                   stride=2,
        │   │   │   │                   padding=0,
        │   │   │   │                   dilation=1,
        │   │   │   │                   groups=1,
        │   │   │   │                   bias=True,
        │   │   │   │                   depth_index=(),
        │   │   │   │                   device=None,
        │   │   │   │                   dtype=None,
        │   │   │   │                   init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │                   init_bias=torch.nn.init.zeros_
        │   │   │   │               )
        │   │   │   │           ],
        │   │   │   │           repeat=1,
        │   │   │   │           depth_index=()
        │   │   │   │       ),
        │   │   │   │       block=[
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=768,
        │   │   │   │       │       channels_out=768,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=768,
        │   │   │   │       │           channels_out=768,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=768,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=768,
        │   │   │   │       │           channels_out=768,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=768,
        │   │   │   │       │           channels_out=768,
        │   │   │   │       │           channels_hidden=3_072,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.09428571909666061
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=768,
        │   │   │   │       │       channels_out=768,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=768,
        │   │   │   │       │           channels_out=768,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=768,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=768,
        │   │   │   │       │           channels_out=768,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=768,
        │   │   │   │       │           channels_out=768,
        │   │   │   │       │           channels_hidden=3_072,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.09714286029338837
        │   │   │   │       │   ),
        │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │       │       channels_in=768,
        │   │   │   │       │       channels_out=768,
        │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │       │           channels_in=768,
        │   │   │   │       │           channels_out=768,
        │   │   │   │       │           kernel_size=7,
        │   │   │   │       │           stride=1,
        │   │   │   │       │           padding=3,
        │   │   │   │       │           dilation=1,
        │   │   │   │       │           groups=768,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None,
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │       │       ),
        │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │       │           channels_in=768,
        │   │   │   │       │           channels_out=768,
        │   │   │   │       │           eps=1e-06,
        │   │   │   │       │           elementwise_affine=True,
        │   │   │   │       │           device=None,
        │   │   │   │       │           dtype=None
        │   │   │   │       │       ),
        │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │       │           channels_in=768,
        │   │   │   │       │           channels_out=768,
        │   │   │   │       │           channels_hidden=3_072,
        │   │   │   │       │           expansion=4.0,
        │   │   │   │       │           round_to=1,
        │   │   │   │       │           gate=False,
        │   │   │   │       │           bias=True,
        │   │   │   │       │           split_gate_projection=False,
        │   │   │   │       │           norm=None,
        │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │       │           depth_index=(),
        │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │   │   │       │           up_proj=None,
        │   │   │   │       │           middle=None,
        │   │   │   │       │           down_proj=None,
        │   │   │   │       │           shard=None
        │   │   │   │       │       ),
        │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │       │       drop_path=0.10000000149011612
        │   │   │   │       │   )
        │   │   │   │       ]
        │   │   │   │   )
        │   │   │   ],
        │   │   │   head=Identity.Config(channels_in=768, channels_out=768),
        │   │   │   drop_path=0.1,
        │   │   │   drop_path_schedule=priml.model.convnext.linear_drop_path
        │   │   ),
        │   │   proj_in=Linear.Config(
        │   │       channels_in=768,
        │   │       channels_out=384,
        │   │       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
        │   │   ),
        │   │   norm_in=LayerNorm.Config(
        │   │       channels_in=384,
        │   │       channels_out=384,
        │   │       eps=1e-05,
        │   │       elementwise_affine=True,
        │   │       device=None,
        │   │       dtype=None
        │   │   ),
        │   │   slot=Embedding.Config(
        │   │       channels_in=10,
        │   │       channels_out=384,
        │   │       padding_idx=None,
        │   │       device=None,
        │   │       dtype=None,
        │   │       lookup_dtype=None,
        │   │       shard=None,
        │   │       depth_index=(),
        │   │       init_weight=torch.nn.init.zeros_
        │   │   ),
        │   │   view=Embedding.Config(
        │   │       channels_in=6,
        │   │       channels_out=384,
        │   │       padding_idx=None,
        │   │       device=None,
        │   │       dtype=None,
        │   │       lookup_dtype=None,
        │   │       shard=None,
        │   │       depth_index=(),
        │   │       init_weight=torch.nn.init.zeros_
        │   │   ),
        │   │   geom=Linear.Config(
        │   │       channels_in=4,
        │   │       channels_out=384,
        │   │       bias=False,
        │   │       depth_index=(),
        │   │       device=None,
        │   │       dtype=None,
        │   │       init_weight=torch.nn.init.zeros_,
        │   │       init_bias=torch.nn.init.zeros_,
        │   │       linear_fn=<built-in function linear>,
        │   │       shard=None
        │   │   ),
        │   │   dropout=Dropout.Config(
        │   │       channels_in=384,
        │   │       channels_out=384,
        │   │       p=0.1,
        │   │       inplace=False
        │   │   ),
        │   │   win_proj=Linear.Config(
        │   │       channels_in=384,
        │   │       channels_out=384,
        │   │       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
        │   │   ),
        │   │   norm_out=LayerNorm.Config(
        │   │       channels_in=384,
        │   │       channels_out=384,
        │   │       eps=1e-05,
        │   │       elementwise_affine=True,
        │   │       device=None,
        │   │       dtype=None
        │   │   ),
        │   │   norm_spatial=LayerNorm.Config(
        │   │       channels_in=768,
        │   │       channels_out=768,
        │   │       eps=1e-05,
        │   │       elementwise_affine=False,
        │   │       device=None,
        │   │       dtype=None
        │   │   ),
        │   │   norm_sp_head=LayerNorm.Config(
        │   │       channels_in=768,
        │   │       channels_out=768,
        │   │       eps=1e-06,
        │   │       elementwise_affine=True,
        │   │       device=None,
        │   │       dtype=None
        │   │   ),
        │   │   mean=(0.485, 0.456, 0.406),
        │   │   std=(0.229, 0.224, 0.225),
        │   │   pixel_scale=255.0,
        │   │   pad_multiple=32,
        │   │   init_weight=functools.partial(priml.model.init.normal, std=0.02),
        │   │   init_bias_extent=functools.partial(torch.nn.init.constant_, val=-4.0),
        │   │   init_prior_state=functools.partial(torch.nn.init.constant_, val=0.5),
        │   │   sp_key=Linear.Config(
        │   │       channels_in=768,
        │   │       channels_out=128,
        │   │       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
        │   │   ),
        │   │   sp_pos=Linear.Config(
        │   │       channels_in=16,
        │   │       channels_out=128,
        │   │       bias=False,
        │   │       depth_index=(),
        │   │       device=None,
        │   │       dtype=None,
        │   │       init_weight=torch.nn.init.zeros_,
        │   │       init_bias=torch.nn.init.zeros_,
        │   │       linear_fn=<built-in function linear>,
        │   │       shard=None
        │   │   ),
        │   │   frequencies=(1.0, 2.0, 4.0, 8.0),
        │   │   state_counts=(6, 6, 5, 5, 6, 6, 6, 6, 5, 6, 6, 8),
        │   │   binary_parents=(4, 5, 6, 4, 5, 6, 4, 5, 6, 1, 11, 10, 0),
        │   │   extent_classes=(4, 5, 6)
        │   ),
        │   anatomy_checkpoint=ReleasedState.Config(
        │       base_dir='/opt/scratch',
        │       file=LocalFile.Config(
        │           path=PosixPath('/opt/scratch/datasets/rsna-knee-orthofoundation-v1/anatomy.pt'),
        │           sha256='5904dd16eb9a4fd7a1caf90ac977f8fe2ef364908598c23c209db3fe7fd1563e'
        │       ),
        │       key='state_dict',
        │       weights_only=True
        │   ),
        │   anatomy=ReleasedVNet.Config(
        │   │   channels_in=1,
        │   │   channels_out=4,
        │   │   channels_hidden=(16, 32, 64, 128, 256),
        │   │   input=VNetInput.Config(
        │   │       channels_in=1,
        │   │       channels_out=16,
        │   │       layer=ConvNorm.Config(
        │   │           channels_in=1,
        │   │           channels_out=16,
        │   │           conv=Conv2d.Config(
        │   │               channels_in=1,
        │   │               channels_out=16,
        │   │               kernel_size=5,
        │   │               stride=1,
        │   │               padding=2,
        │   │               dilation=1,
        │   │               groups=1,
        │   │               bias=True,
        │   │               depth_index=(),
        │   │               device=None,
        │   │               dtype=None,
        │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │               init_bias=torch.nn.init.zeros_
        │   │           ),
        │   │           norm=BatchNorm2d.Config(
        │   │               channels_in=16,
        │   │               channels_out=16,
        │   │               momentum=0.1,
        │   │               eps=1e-05,
        │   │               elementwise_affine=True,
        │   │               device=None,
        │   │               dtype=None
        │   │           )
        │   │       ),
        │   │       activation=priml.model.swiglu.elu
        │   │   ),
        │   │   down=[
        │   │   │   VNetDown.Config(
        │   │   │       channels_in=16,
        │   │   │       channels_out=32,
        │   │   │       num_layers=1,
        │   │   │       proj=ConvNorm.Config(
        │   │   │           channels_in=16,
        │   │   │           channels_out=32,
        │   │   │           conv=Conv2d.Config(
        │   │   │               channels_in=16,
        │   │   │               channels_out=32,
        │   │   │               kernel_size=2,
        │   │   │               stride=2,
        │   │   │               padding=0,
        │   │   │               dilation=1,
        │   │   │               groups=1,
        │   │   │               bias=True,
        │   │   │               depth_index=(),
        │   │   │               device=None,
        │   │   │               dtype=None,
        │   │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │               init_bias=torch.nn.init.zeros_
        │   │   │           ),
        │   │   │           norm=BatchNorm2d.Config(
        │   │   │               channels_in=32,
        │   │   │               channels_out=32,
        │   │   │               momentum=0.1,
        │   │   │               eps=1e-05,
        │   │   │               elementwise_affine=True,
        │   │   │               device=None,
        │   │   │               dtype=None
        │   │   │           )
        │   │   │       ),
        │   │   │       block=ConvNorm.Config(
        │   │   │           channels_in=32,
        │   │   │           channels_out=32,
        │   │   │           conv=Conv2d.Config(
        │   │   │               channels_in=32,
        │   │   │               channels_out=32,
        │   │   │               kernel_size=5,
        │   │   │               stride=1,
        │   │   │               padding=2,
        │   │   │               dilation=1,
        │   │   │               groups=1,
        │   │   │               bias=True,
        │   │   │               depth_index=(),
        │   │   │               device=None,
        │   │   │               dtype=None,
        │   │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │               init_bias=torch.nn.init.zeros_
        │   │   │           ),
        │   │   │           norm=BatchNorm2d.Config(
        │   │   │               channels_in=32,
        │   │   │               channels_out=32,
        │   │   │               momentum=0.1,
        │   │   │               eps=1e-05,
        │   │   │               elementwise_affine=True,
        │   │   │               device=None,
        │   │   │               dtype=None
        │   │   │           )
        │   │   │       ),
        │   │   │       dropout=None,
        │   │   │       activation=priml.model.swiglu.elu
        │   │   │   ),
        │   │   │   VNetDown.Config(
        │   │   │       channels_in=32,
        │   │   │       channels_out=64,
        │   │   │       num_layers=2,
        │   │   │       proj=ConvNorm.Config(
        │   │   │           channels_in=32,
        │   │   │           channels_out=64,
        │   │   │           conv=Conv2d.Config(
        │   │   │               channels_in=32,
        │   │   │               channels_out=64,
        │   │   │               kernel_size=2,
        │   │   │               stride=2,
        │   │   │               padding=0,
        │   │   │               dilation=1,
        │   │   │               groups=1,
        │   │   │               bias=True,
        │   │   │               depth_index=(),
        │   │   │               device=None,
        │   │   │               dtype=None,
        │   │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │               init_bias=torch.nn.init.zeros_
        │   │   │           ),
        │   │   │           norm=BatchNorm2d.Config(
        │   │   │               channels_in=64,
        │   │   │               channels_out=64,
        │   │   │               momentum=0.1,
        │   │   │               eps=1e-05,
        │   │   │               elementwise_affine=True,
        │   │   │               device=None,
        │   │   │               dtype=None
        │   │   │           )
        │   │   │       ),
        │   │   │       block=ConvNorm.Config(
        │   │   │           channels_in=64,
        │   │   │           channels_out=64,
        │   │   │           conv=Conv2d.Config(
        │   │   │               channels_in=64,
        │   │   │               channels_out=64,
        │   │   │               kernel_size=5,
        │   │   │               stride=1,
        │   │   │               padding=2,
        │   │   │               dilation=1,
        │   │   │               groups=1,
        │   │   │               bias=True,
        │   │   │               depth_index=(),
        │   │   │               device=None,
        │   │   │               dtype=None,
        │   │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │               init_bias=torch.nn.init.zeros_
        │   │   │           ),
        │   │   │           norm=BatchNorm2d.Config(
        │   │   │               channels_in=64,
        │   │   │               channels_out=64,
        │   │   │               momentum=0.1,
        │   │   │               eps=1e-05,
        │   │   │               elementwise_affine=True,
        │   │   │               device=None,
        │   │   │               dtype=None
        │   │   │           )
        │   │   │       ),
        │   │   │       dropout=None,
        │   │   │       activation=priml.model.swiglu.elu
        │   │   │   ),
        │   │   │   VNetDown.Config(
        │   │   │       channels_in=64,
        │   │   │       channels_out=128,
        │   │   │       num_layers=3,
        │   │   │       proj=ConvNorm.Config(
        │   │   │           channels_in=64,
        │   │   │           channels_out=128,
        │   │   │           conv=Conv2d.Config(
        │   │   │               channels_in=64,
        │   │   │               channels_out=128,
        │   │   │               kernel_size=2,
        │   │   │               stride=2,
        │   │   │               padding=0,
        │   │   │               dilation=1,
        │   │   │               groups=1,
        │   │   │               bias=True,
        │   │   │               depth_index=(),
        │   │   │               device=None,
        │   │   │               dtype=None,
        │   │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │               init_bias=torch.nn.init.zeros_
        │   │   │           ),
        │   │   │           norm=BatchNorm2d.Config(
        │   │   │               channels_in=128,
        │   │   │               channels_out=128,
        │   │   │               momentum=0.1,
        │   │   │               eps=1e-05,
        │   │   │               elementwise_affine=True,
        │   │   │               device=None,
        │   │   │               dtype=None
        │   │   │           )
        │   │   │       ),
        │   │   │       block=ConvNorm.Config(
        │   │   │           channels_in=128,
        │   │   │           channels_out=128,
        │   │   │           conv=Conv2d.Config(
        │   │   │               channels_in=128,
        │   │   │               channels_out=128,
        │   │   │               kernel_size=5,
        │   │   │               stride=1,
        │   │   │               padding=2,
        │   │   │               dilation=1,
        │   │   │               groups=1,
        │   │   │               bias=True,
        │   │   │               depth_index=(),
        │   │   │               device=None,
        │   │   │               dtype=None,
        │   │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │               init_bias=torch.nn.init.zeros_
        │   │   │           ),
        │   │   │           norm=BatchNorm2d.Config(
        │   │   │               channels_in=128,
        │   │   │               channels_out=128,
        │   │   │               momentum=0.1,
        │   │   │               eps=1e-05,
        │   │   │               elementwise_affine=True,
        │   │   │               device=None,
        │   │   │               dtype=None
        │   │   │           )
        │   │   │       ),
        │   │   │       dropout=Dropout2d.Config(
        │   │   │           channels_in=128,
        │   │   │           channels_out=128,
        │   │   │           p=0.0,
        │   │   │           inplace=False
        │   │   │       ),
        │   │   │       activation=priml.model.swiglu.elu
        │   │   │   ),
        │   │   │   VNetDown.Config(
        │   │   │       channels_in=128,
        │   │   │       channels_out=256,
        │   │   │       num_layers=2,
        │   │   │       proj=ConvNorm.Config(
        │   │   │           channels_in=128,
        │   │   │           channels_out=256,
        │   │   │           conv=Conv2d.Config(
        │   │   │               channels_in=128,
        │   │   │               channels_out=256,
        │   │   │               kernel_size=2,
        │   │   │               stride=2,
        │   │   │               padding=0,
        │   │   │               dilation=1,
        │   │   │               groups=1,
        │   │   │               bias=True,
        │   │   │               depth_index=(),
        │   │   │               device=None,
        │   │   │               dtype=None,
        │   │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │               init_bias=torch.nn.init.zeros_
        │   │   │           ),
        │   │   │           norm=BatchNorm2d.Config(
        │   │   │               channels_in=256,
        │   │   │               channels_out=256,
        │   │   │               momentum=0.1,
        │   │   │               eps=1e-05,
        │   │   │               elementwise_affine=True,
        │   │   │               device=None,
        │   │   │               dtype=None
        │   │   │           )
        │   │   │       ),
        │   │   │       block=ConvNorm.Config(
        │   │   │           channels_in=256,
        │   │   │           channels_out=256,
        │   │   │           conv=Conv2d.Config(
        │   │   │               channels_in=256,
        │   │   │               channels_out=256,
        │   │   │               kernel_size=5,
        │   │   │               stride=1,
        │   │   │               padding=2,
        │   │   │               dilation=1,
        │   │   │               groups=1,
        │   │   │               bias=True,
        │   │   │               depth_index=(),
        │   │   │               device=None,
        │   │   │               dtype=None,
        │   │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │               init_bias=torch.nn.init.zeros_
        │   │   │           ),
        │   │   │           norm=BatchNorm2d.Config(
        │   │   │               channels_in=256,
        │   │   │               channels_out=256,
        │   │   │               momentum=0.1,
        │   │   │               eps=1e-05,
        │   │   │               elementwise_affine=True,
        │   │   │               device=None,
        │   │   │               dtype=None
        │   │   │           )
        │   │   │       ),
        │   │   │       dropout=Dropout2d.Config(
        │   │   │           channels_in=256,
        │   │   │           channels_out=256,
        │   │   │           p=0.0,
        │   │   │           inplace=False
        │   │   │       ),
        │   │   │       activation=priml.model.swiglu.elu
        │   │   │   )
        │   │   ],
        │   │   up=[
        │   │   │   VNetUp.Config(
        │   │   │       channels_in=256,
        │   │   │       channels_skip=128,
        │   │   │       num_layers=2,
        │   │   │       proj=ConvTranspose2d.Config(
        │   │   │           channels_in=256,
        │   │   │           channels_out=128,
        │   │   │           kernel_size=2,
        │   │   │           stride=2,
        │   │   │           padding=0,
        │   │   │           output_padding=0,
        │   │   │           dilation=1,
        │   │   │           groups=1,
        │   │   │           bias=True,
        │   │   │           depth_index=(),
        │   │   │           device=None,
        │   │   │           dtype=None,
        │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │           init_bias=torch.nn.init.zeros_
        │   │   │       ),
        │   │   │       norm=BatchNorm2d.Config(
        │   │   │           channels_in=128,
        │   │   │           channels_out=128,
        │   │   │           momentum=0.1,
        │   │   │           eps=1e-05,
        │   │   │           elementwise_affine=True,
        │   │   │           device=None,
        │   │   │           dtype=None
        │   │   │       ),
        │   │   │       block=ConvNorm.Config(
        │   │   │           channels_in=256,
        │   │   │           channels_out=256,
        │   │   │           conv=Conv2d.Config(
        │   │   │               channels_in=256,
        │   │   │               channels_out=256,
        │   │   │               kernel_size=5,
        │   │   │               stride=1,
        │   │   │               padding=2,
        │   │   │               dilation=1,
        │   │   │               groups=1,
        │   │   │               bias=True,
        │   │   │               depth_index=(),
        │   │   │               device=None,
        │   │   │               dtype=None,
        │   │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │               init_bias=torch.nn.init.zeros_
        │   │   │           ),
        │   │   │           norm=BatchNorm2d.Config(
        │   │   │               channels_in=256,
        │   │   │               channels_out=256,
        │   │   │               momentum=0.1,
        │   │   │               eps=1e-05,
        │   │   │               elementwise_affine=True,
        │   │   │               device=None,
        │   │   │               dtype=None
        │   │   │           )
        │   │   │       ),
        │   │   │       dropout=Dropout2d.Config(
        │   │   │           channels_in=256,
        │   │   │           channels_out=256,
        │   │   │           p=0.0,
        │   │   │           inplace=False
        │   │   │       ),
        │   │   │       dropout_skip=Dropout2d.Config(
        │   │   │           channels_in=128,
        │   │   │           channels_out=128,
        │   │   │           p=0.5,
        │   │   │           inplace=False
        │   │   │       ),
        │   │   │       activation=priml.model.swiglu.elu
        │   │   │   ),
        │   │   │   VNetUp.Config(
        │   │   │       channels_in=256,
        │   │   │       channels_skip=64,
        │   │   │       num_layers=2,
        │   │   │       proj=ConvTranspose2d.Config(
        │   │   │           channels_in=256,
        │   │   │           channels_out=64,
        │   │   │           kernel_size=2,
        │   │   │           stride=2,
        │   │   │           padding=0,
        │   │   │           output_padding=0,
        │   │   │           dilation=1,
        │   │   │           groups=1,
        │   │   │           bias=True,
        │   │   │           depth_index=(),
        │   │   │           device=None,
        │   │   │           dtype=None,
        │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │           init_bias=torch.nn.init.zeros_
        │   │   │       ),
        │   │   │       norm=BatchNorm2d.Config(
        │   │   │           channels_in=64,
        │   │   │           channels_out=64,
        │   │   │           momentum=0.1,
        │   │   │           eps=1e-05,
        │   │   │           elementwise_affine=True,
        │   │   │           device=None,
        │   │   │           dtype=None
        │   │   │       ),
        │   │   │       block=ConvNorm.Config(
        │   │   │           channels_in=128,
        │   │   │           channels_out=128,
        │   │   │           conv=Conv2d.Config(
        │   │   │               channels_in=128,
        │   │   │               channels_out=128,
        │   │   │               kernel_size=5,
        │   │   │               stride=1,
        │   │   │               padding=2,
        │   │   │               dilation=1,
        │   │   │               groups=1,
        │   │   │               bias=True,
        │   │   │               depth_index=(),
        │   │   │               device=None,
        │   │   │               dtype=None,
        │   │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │               init_bias=torch.nn.init.zeros_
        │   │   │           ),
        │   │   │           norm=BatchNorm2d.Config(
        │   │   │               channels_in=128,
        │   │   │               channels_out=128,
        │   │   │               momentum=0.1,
        │   │   │               eps=1e-05,
        │   │   │               elementwise_affine=True,
        │   │   │               device=None,
        │   │   │               dtype=None
        │   │   │           )
        │   │   │       ),
        │   │   │       dropout=Dropout2d.Config(
        │   │   │           channels_in=256,
        │   │   │           channels_out=256,
        │   │   │           p=0.0,
        │   │   │           inplace=False
        │   │   │       ),
        │   │   │       dropout_skip=Dropout2d.Config(
        │   │   │           channels_in=64,
        │   │   │           channels_out=64,
        │   │   │           p=0.5,
        │   │   │           inplace=False
        │   │   │       ),
        │   │   │       activation=priml.model.swiglu.elu
        │   │   │   ),
        │   │   │   VNetUp.Config(
        │   │   │       channels_in=128,
        │   │   │       channels_skip=32,
        │   │   │       num_layers=1,
        │   │   │       proj=ConvTranspose2d.Config(
        │   │   │           channels_in=128,
        │   │   │           channels_out=32,
        │   │   │           kernel_size=2,
        │   │   │           stride=2,
        │   │   │           padding=0,
        │   │   │           output_padding=0,
        │   │   │           dilation=1,
        │   │   │           groups=1,
        │   │   │           bias=True,
        │   │   │           depth_index=(),
        │   │   │           device=None,
        │   │   │           dtype=None,
        │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │           init_bias=torch.nn.init.zeros_
        │   │   │       ),
        │   │   │       norm=BatchNorm2d.Config(
        │   │   │           channels_in=32,
        │   │   │           channels_out=32,
        │   │   │           momentum=0.1,
        │   │   │           eps=1e-05,
        │   │   │           elementwise_affine=True,
        │   │   │           device=None,
        │   │   │           dtype=None
        │   │   │       ),
        │   │   │       block=ConvNorm.Config(
        │   │   │           channels_in=64,
        │   │   │           channels_out=64,
        │   │   │           conv=Conv2d.Config(
        │   │   │               channels_in=64,
        │   │   │               channels_out=64,
        │   │   │               kernel_size=5,
        │   │   │               stride=1,
        │   │   │               padding=2,
        │   │   │               dilation=1,
        │   │   │               groups=1,
        │   │   │               bias=True,
        │   │   │               depth_index=(),
        │   │   │               device=None,
        │   │   │               dtype=None,
        │   │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │               init_bias=torch.nn.init.zeros_
        │   │   │           ),
        │   │   │           norm=BatchNorm2d.Config(
        │   │   │               channels_in=64,
        │   │   │               channels_out=64,
        │   │   │               momentum=0.1,
        │   │   │               eps=1e-05,
        │   │   │               elementwise_affine=True,
        │   │   │               device=None,
        │   │   │               dtype=None
        │   │   │           )
        │   │   │       ),
        │   │   │       dropout=None,
        │   │   │       dropout_skip=Dropout2d.Config(
        │   │   │           channels_in=32,
        │   │   │           channels_out=32,
        │   │   │           p=0.5,
        │   │   │           inplace=False
        │   │   │       ),
        │   │   │       activation=priml.model.swiglu.elu
        │   │   │   ),
        │   │   │   VNetUp.Config(
        │   │   │       channels_in=64,
        │   │   │       channels_skip=16,
        │   │   │       num_layers=1,
        │   │   │       proj=ConvTranspose2d.Config(
        │   │   │           channels_in=64,
        │   │   │           channels_out=16,
        │   │   │           kernel_size=2,
        │   │   │           stride=2,
        │   │   │           padding=0,
        │   │   │           output_padding=0,
        │   │   │           dilation=1,
        │   │   │           groups=1,
        │   │   │           bias=True,
        │   │   │           depth_index=(),
        │   │   │           device=None,
        │   │   │           dtype=None,
        │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │           init_bias=torch.nn.init.zeros_
        │   │   │       ),
        │   │   │       norm=BatchNorm2d.Config(
        │   │   │           channels_in=16,
        │   │   │           channels_out=16,
        │   │   │           momentum=0.1,
        │   │   │           eps=1e-05,
        │   │   │           elementwise_affine=True,
        │   │   │           device=None,
        │   │   │           dtype=None
        │   │   │       ),
        │   │   │       block=ConvNorm.Config(
        │   │   │           channels_in=32,
        │   │   │           channels_out=32,
        │   │   │           conv=Conv2d.Config(
        │   │   │               channels_in=32,
        │   │   │               channels_out=32,
        │   │   │               kernel_size=5,
        │   │   │               stride=1,
        │   │   │               padding=2,
        │   │   │               dilation=1,
        │   │   │               groups=1,
        │   │   │               bias=True,
        │   │   │               depth_index=(),
        │   │   │               device=None,
        │   │   │               dtype=None,
        │   │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │               init_bias=torch.nn.init.zeros_
        │   │   │           ),
        │   │   │           norm=BatchNorm2d.Config(
        │   │   │               channels_in=32,
        │   │   │               channels_out=32,
        │   │   │               momentum=0.1,
        │   │   │               eps=1e-05,
        │   │   │               elementwise_affine=True,
        │   │   │               device=None,
        │   │   │               dtype=None
        │   │   │           )
        │   │   │       ),
        │   │   │       dropout=None,
        │   │   │       dropout_skip=Dropout2d.Config(
        │   │   │           channels_in=16,
        │   │   │           channels_out=16,
        │   │   │           p=0.5,
        │   │   │           inplace=False
        │   │   │       ),
        │   │   │       activation=priml.model.swiglu.elu
        │   │   │   )
        │   │   ],
        │   │   output=VNetOutput.Config(
        │   │       channels_in=32,
        │   │       channels_out=4,
        │   │       conv=ConvNorm.Config(
        │   │           channels_in=32,
        │   │           channels_out=4,
        │   │           conv=Conv2d.Config(
        │   │               channels_in=32,
        │   │               channels_out=4,
        │   │               kernel_size=5,
        │   │               stride=1,
        │   │               padding=2,
        │   │               dilation=1,
        │   │               groups=1,
        │   │               bias=True,
        │   │               depth_index=(),
        │   │               device=None,
        │   │               dtype=None,
        │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │               init_bias=torch.nn.init.zeros_
        │   │           ),
        │   │           norm=BatchNorm2d.Config(
        │   │               channels_in=4,
        │   │               channels_out=4,
        │   │               momentum=0.1,
        │   │               eps=1e-05,
        │   │               elementwise_affine=True,
        │   │               device=None,
        │   │               dtype=None
        │   │           )
        │   │       ),
        │   │       proj=Conv2d.Config(
        │   │           channels_in=4,
        │   │           channels_out=4,
        │   │           kernel_size=1,
        │   │           stride=1,
        │   │           padding=0,
        │   │           dilation=1,
        │   │           groups=1,
        │   │           bias=True,
        │   │           depth_index=(),
        │   │           device=None,
        │   │           dtype=None,
        │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │           init_bias=torch.nn.init.zeros_
        │   │       ),
        │   │       activation=priml.model.swiglu.elu
        │   │   )
        │   ),
        │   side_checkpoint=ReleasedState.Config(
        │       base_dir='/opt/scratch',
        │       file=LocalFile.Config(
        │           path=PosixPath('/opt/scratch/datasets/rsna-knee-r7-v2/side_cls.pt'),
        │           sha256='ab623e2b619db4917639220fed15202e8c8aedd07c8e298a66a2afc0f153650a'
        │       ),
        │       key=None,
        │       weights_only=True
        │   ),
        │   side=ResNet.Config(
        │   │   channels_in=1,
        │   │   channels_out=1,
        │   │   stem=Conv2d.Config(
        │   │       channels_in=1,
        │   │       channels_out=64,
        │   │       kernel_size=7,
        │   │       stride=2,
        │   │       padding=3,
        │   │       dilation=1,
        │   │       groups=1,
        │   │       bias=False,
        │   │       depth_index=(),
        │   │       device=None,
        │   │       dtype=None,
        │   │       init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │       init_bias=torch.nn.init.zeros_
        │   │   ),
        │   │   norm=BatchNorm2d.Config(
        │   │       channels_in=64,
        │   │       channels_out=64,
        │   │       momentum=0.1,
        │   │       eps=1e-05,
        │   │       elementwise_affine=True,
        │   │       device=None,
        │   │       dtype=None
        │   │   ),
        │   │   channels_hidden=(64, 128, 256, 512),
        │   │   blocks_per_stage=(3, 4, 6, 3),
        │   │   stage=[
        │   │   │   ResNetStage.Config(
        │   │   │       channels_in=64,
        │   │   │       channels_out=64,
        │   │   │       stride=1,
        │   │   │       blocks_per_stage=3,
        │   │   │       block=[
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=64,
        │   │   │       │       channels_out=64,
        │   │   │       │       stride=1,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=64,
        │   │   │       │       channels_out=64,
        │   │   │       │       stride=1,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=64,
        │   │   │       │       channels_out=64,
        │   │   │       │       stride=1,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=64,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   )
        │   │   │       ]
        │   │   │   ),
        │   │   │   ResNetStage.Config(
        │   │   │       channels_in=64,
        │   │   │       channels_out=128,
        │   │   │       stride=2,
        │   │   │       blocks_per_stage=4,
        │   │   │       block=[
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=64,
        │   │   │       │       channels_out=128,
        │   │   │       │       stride=2,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=2,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=64,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=2,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=128,
        │   │   │       │       channels_out=128,
        │   │   │       │       stride=1,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=128,
        │   │   │       │       channels_out=128,
        │   │   │       │       stride=1,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=128,
        │   │   │       │       channels_out=128,
        │   │   │       │       stride=1,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=128,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   )
        │   │   │       ]
        │   │   │   ),
        │   │   │   ResNetStage.Config(
        │   │   │       channels_in=128,
        │   │   │       channels_out=256,
        │   │   │       stride=2,
        │   │   │       blocks_per_stage=6,
        │   │   │       block=[
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=128,
        │   │   │       │       channels_out=256,
        │   │   │       │       stride=2,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=2,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=128,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=2,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=256,
        │   │   │       │       channels_out=256,
        │   │   │       │       stride=1,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=256,
        │   │   │       │       channels_out=256,
        │   │   │       │       stride=1,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=256,
        │   │   │       │       channels_out=256,
        │   │   │       │       stride=1,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=256,
        │   │   │       │       channels_out=256,
        │   │   │       │       stride=1,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=256,
        │   │   │       │       channels_out=256,
        │   │   │       │       stride=1,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=256,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   )
        │   │   │       ]
        │   │   │   ),
        │   │   │   ResNetStage.Config(
        │   │   │       channels_in=256,
        │   │   │       channels_out=512,
        │   │   │       stride=2,
        │   │   │       blocks_per_stage=3,
        │   │   │       block=[
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=256,
        │   │   │       │       channels_out=512,
        │   │   │       │       stride=2,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=512,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=2,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=256,
        │   │   │       │           channels_out=512,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=2,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       stride=1,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   ),
        │   │   │       │   ResNetBlock.Config(
        │   │   │       │       channels_in=512,
        │   │   │       │       channels_out=512,
        │   │   │       │       stride=1,
        │   │   │       │       conv1=Conv2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm1=BatchNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       conv2=Conv2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           kernel_size=3,
        │   │   │       │           stride=1,
        │   │   │       │           padding=1,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm2=BatchNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       projection=Conv2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           kernel_size=1,
        │   │   │       │           stride=1,
        │   │   │       │           padding=0,
        │   │   │       │           dilation=1,
        │   │   │       │           groups=1,
        │   │   │       │           bias=False,
        │   │   │       │           depth_index=(),
        │   │   │       │           device=None,
        │   │   │       │           dtype=None,
        │   │   │       │           init_weight=functools.partial(torch.nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'),
        │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │       │       ),
        │   │   │       │       norm_projection=BatchNorm2d.Config(
        │   │   │       │           channels_in=512,
        │   │   │       │           channels_out=512,
        │   │   │       │           momentum=0.1,
        │   │   │       │           eps=1e-05,
        │   │   │       │           elementwise_affine=True,
        │   │   │       │           device=None,
        │   │   │       │           dtype=None
        │   │   │       │       ),
        │   │   │       │       activation=priml.model.resnet.relu_inplace,
        │   │   │       │       init_scale_residual=torch.nn.init.zeros_
        │   │   │       │   )
        │   │   │       ]
        │   │   │   )
        │   │   ],
        │   │   proj_out=Linear.Config(
        │   │       channels_in=512,
        │   │       channels_out=1,
        │   │       bias=True,
        │   │       depth_index=(),
        │   │       device=None,
        │   │       dtype=None,
        │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │       init_bias=torch.nn.init.zeros_,
        │   │       linear_fn=<built-in function linear>,
        │   │       shard=None
        │   │   ),
        │   │   activation=priml.model.resnet.relu_inplace
        │   ),
        │   preparation=StudyPreparation.Config(
        │   │   eval_stride=1,
        │   │   side=SideCheck.Config(
        │   │       image_size=160,
        │   │       fractions=[0.4, 0.5, 0.6],
        │   │       threshold=2.0,
        │   │       max_x=inf,
        │   │       position_mm=20.0,
        │   │       window_percentiles=(1.0, 99.0),
        │   │       percentile_stride=4,
        │   │       mean=0.45,
        │   │       std=0.25
        │   │   ),
        │   │   anatomy=AnatomySampler.Config(
        │   │       slices=7,
        │   │       span_mm=4.0,
        │   │       pixel_mm=0.3125,
        │   │       pad_multiple=16,
        │   │       oblique_cos=0.85
        │   │   ),
        │   │   anchors=AnchorConsensus.Config(
        │   │       agreement_mm=20.0,
        │   │       prior_mm=35.0,
        │   │       mass_min=250,
        │   │       tissue_threshold=0.5,
        │   │       tissue_min_pixels=40,
        │   │       patella_slice_min_pixels=60,
        │   │       patella_mass_min=120,
        │   │       patella_margin_mm=10.0,
        │   │       field_min=0.15,
        │   │       field_max=0.85
        │   │   ),
        │   │   render=R7Render.Config(
        │   │       image_size=384,
        │   │       local_mm=100.0,
        │   │       thin_mm=1.75,
        │   │       slab_mm=3.5,
        │   │       view_offsets_mm={   'Sagittal': [
        │   │               (1, 0.0, 0.0),
        │   │               (4, -25.0, 40.0),
        │   │               (5, 5.0, -35.0)
        │   │           ],
        │   │           'Coronal': [(2, 0.0, -30.0), (3, 0.0, 30.0)],
        │   │           'Axial': [
        │   │               (1, 0.0, 0.0),
        │   │               (4, -40.0, 0.0),
        │   │               (5, 35.0, 0.0)
        │   │           ]},
        │   │       display_floor=0.02,
        │   │       display_knee=0.9
        │   │   ),
        │   │   fluid=FluidMeasure.Config(
        │   │       threshold=229,
        │   │       opening_mm=1.5,
        │   │       foreground_min=10
        │   │   )
        │   ),
        │   score=R7Score.Config(
        │       device='cuda',
        │       crop=352,
        │       chunk_size=192,
        │       dtype_autocast=torch.bfloat16
        │   ),
        │   postprocess=R7Postprocess.Config(
        │       base_dir='/opt/scratch',
        │       fluid_file=LocalFile.Config(
        │           path=PosixPath('/opt/scratch/datasets/rsna-knee-r7-v2/fluid_stacker.json'),
        │           sha256='be5a1664077300db1fdd0af8daaf44a3ebc41f91152e4f0ec124f805a7a51efb'
        │       ),
        │       state_file=LocalFile.Config(
        │           path=PosixPath('/opt/scratch/datasets/rsna-knee-r7-v2/state_stacker.json'),
        │           sha256='7b54a3f774dad3f842d68aaccbb03771806c26c0a3ca3a07fea8575a21f3ead3'
        │       ),
        │       spec_file=LocalFile.Config(
        │           path=PosixPath('/opt/scratch/datasets/rsna-knee-r7-v2/state_spec.json'),
        │           sha256='2c976e70773b745cb357b992ba04946c157aa478e65d82f2883039c038b6ab3c'
        │       )
        │   ),
        │   execution=EncoderExecution.Config(
        │       mode='default',
        │       backend='inductor',
        │       dynamic=False,
        │       probability_tolerance=0.001,
        │       synthetic_tolerance=0.05,
        │       seed=20_261_010,
        │       recompile_limit=64,
        │       accumulated_recompile_limit=2_048
        │   ),
        │   cudnn_benchmark=False,
        │   allow_tf32_cudnn=True,
        │   allow_tf32_matmul=False
        ),
        blend=RankBlend.Config(
            groups=[
                RankGroup.Config(
                    weight=0.5,
                    members=[
                        RankMember.Config(
                            name='Raptor',
                            weight=0.65,
                            max_failure_fraction=0.03
                        ),
                        RankMember.Config(
                            name='Raptor-224',
                            weight=0.35,
                            max_failure_fraction=0.03
                        )
                    ]
                ),
                RankGroup.Config(
                    weight=0.5,
                    members=[
                        RankMember.Config(
                            name='R7',
                            weight=1.0,
                            max_failure_fraction=0.5
                        )
                    ]
                )
            ]
        ),
        budget=WallBudget.Config(
            seconds=32_400,
            reserve_evaluation_sec=2519.9999999999964,
            clock=<built-in function perf_counter>,
            receipt_path=None
        )
    )
