ColdRun.Config(
        loop=KneeTrainLoop.Config(
        │   study_name='rsna_knee_training',
        │   experiment_name='exp000',
        │   base_dir='/opt/scratch',
        │   working_dir=PosixPath('/opt/scratch/runs/rsna_knee_training/exp000'),
        │   doc='',
        │   step=TrainStep.Config(
        │   │   model=Raptor.Config(
        │   │   │   backbone=ImageNetEncoder.Config(
        │   │   │   │   backbone=TinyConvNeXt.Config(
        │   │   │   │   │   channels_in=3,
        │   │   │   │   │   channels_hidden=(96, 192, 384, 768),
        │   │   │   │   │   blocks_per_stage=(3, 3, 9, 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=None
        │   │   │   │   │       ),
        │   │   │   │   │       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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=96,
        │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │       │               channels_out=96,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=96,
        │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │       │               channels_out=96,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.0058823529411764705
        │   │   │   │   │   │       │   ),
        │   │   │   │   │   │       │   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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=96,
        │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │       │               channels_out=96,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.011764705882352941
        │   │   │   │   │   │       │   )
        │   │   │   │   │   │       ]
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   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=None
        │   │   │   │   │   │               )
        │   │   │   │   │   │           ],
        │   │   │   │   │   │           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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=192,
        │   │   │   │   │   │       │               channels_out=768,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=768,
        │   │   │   │   │   │       │               channels_out=192,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.017647058823529415
        │   │   │   │   │   │       │   ),
        │   │   │   │   │   │       │   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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=192,
        │   │   │   │   │   │       │               channels_out=768,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=768,
        │   │   │   │   │   │       │               channels_out=192,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.023529411764705882
        │   │   │   │   │   │       │   ),
        │   │   │   │   │   │       │   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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=192,
        │   │   │   │   │   │       │               channels_out=768,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=768,
        │   │   │   │   │   │       │               channels_out=192,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.029411764705882353
        │   │   │   │   │   │       │   )
        │   │   │   │   │   │       ]
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   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=None
        │   │   │   │   │   │               )
        │   │   │   │   │   │           ],
        │   │   │   │   │   │           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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.03529411764705883
        │   │   │   │   │   │       │   ),
        │   │   │   │   │   │       │   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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.0411764705882353
        │   │   │   │   │   │       │   ),
        │   │   │   │   │   │       │   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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.047058823529411764
        │   │   │   │   │   │       │   ),
        │   │   │   │   │   │       │   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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.052941176470588235
        │   │   │   │   │   │       │   ),
        │   │   │   │   │   │       │   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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.058823529411764705
        │   │   │   │   │   │       │   ),
        │   │   │   │   │   │       │   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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.06470588235294118
        │   │   │   │   │   │       │   ),
        │   │   │   │   │   │       │   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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.07058823529411766
        │   │   │   │   │   │       │   ),
        │   │   │   │   │   │       │   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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.07647058823529412
        │   │   │   │   │   │       │   ),
        │   │   │   │   │   │       │   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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.0823529411764706
        │   │   │   │   │   │       │   )
        │   │   │   │   │   │       ]
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   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=None
        │   │   │   │   │   │               )
        │   │   │   │   │   │           ],
        │   │   │   │   │   │           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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=768,
        │   │   │   │   │   │       │               channels_out=3_072,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=3_072,
        │   │   │   │   │   │       │               channels_out=768,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.08823529411764706
        │   │   │   │   │   │       │   ),
        │   │   │   │   │   │       │   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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=768,
        │   │   │   │   │   │       │               channels_out=3_072,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=3_072,
        │   │   │   │   │   │       │               channels_out=768,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.09411764705882353
        │   │   │   │   │   │       │   ),
        │   │   │   │   │   │       │   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=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       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=False,
        │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=768,
        │   │   │   │   │   │       │               channels_out=3_072,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │       │               channels_in=3_072,
        │   │   │   │   │   │       │               channels_out=768,
        │   │   │   │   │   │       │               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
        │   │   │   │   │   │       │           ),
        │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │       │       ),
        │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │       │       drop_path=0.1
        │   │   │   │   │   │       │   )
        │   │   │   │   │   │       ]
        │   │   │   │   │   │   )
        │   │   │   │   │   ],
        │   │   │   │   │   head=ConvNeXtHead.Config(
        │   │   │   │   │       channels_in=768,
        │   │   │   │   │       norm=LayerNorm.Config(
        │   │   │   │   │           channels_in=768,
        │   │   │   │   │           channels_out=768,
        │   │   │   │   │           eps=1e-06,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       proj_out=Linear.Config(
        │   │   │   │   │           channels_in=768,
        │   │   │   │   │           channels_out=1_000,
        │   │   │   │   │           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
        │   │   │   │   │       )
        │   │   │   │   │   ),
        │   │   │   │   │   drop_path=0.1,
        │   │   │   │   │   drop_path_schedule=priml.baselines.rsna_knee.model.tiny_drop_path,
        │   │   │   │   │   init_weight=functools.partial(torch.nn.init.trunc_normal_, std=0.02),
        │   │   │   │   │   init_bias=torch.nn.init.zeros_
        │   │   │   │   ),
        │   │   │   │   base_dir='/opt/scratch',
        │   │   │   │   working_dir='/datasets/rsna-knee-imagenet',
        │   │   │   │   checkpoint=LocalFile.Config(
        │   │   │   │       path=PosixPath('/opt/scratch/datasets/rsna-knee-imagenet/convnext_tiny-983f1562.pth'),
        │   │   │   │       sha256='983f1562536e84ff750a1576fb08e54de751dbf2e17c0d8a4a13704341fdcd3d'
        │   │   │   │   ),
        │   │   │   │   classifier='head.proj_out'
        │   │   │   ),
        │   │   │   head=FindingAttention.Config(
        │   │   │   │   channels_in=768,
        │   │   │   │   channels_out=12,
        │   │   │   │   channels_hidden=256,
        │   │   │   │   norm=LayerNorm.Config(
        │   │   │   │       channels_in=768,
        │   │   │   │       channels_out=768,
        │   │   │   │       eps=1e-05,
        │   │   │   │       elementwise_affine=True,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None
        │   │   │   │   ),
        │   │   │   │   proj_in=Linear.Config(
        │   │   │   │       channels_in=768,
        │   │   │   │       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.0,
        │   │   │   │       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
        │   │   ),
        │   │   optimizer=CompositeOptimizer.Config(
        │   │       optimizers=[
        │   │           PartialConfig(
        │   │               <class 'torch.optim.adamw.AdamW'>,
        │   │               lr=0.0001,
        │   │               weight_decay=0.01
        │   │           )
        │   │       ],
        │   │       select=[],
        │   │       require_total=True,
        │   │       drop_empty=False
        │   │   ),
        │   │   loss=SimpleLoss.Config(
        │   │       loss_fn=priml.loss.simple_loss.bce_with_logits,
        │   │       target_key='label',
        │   │       kwargs={'reduction': 'none'},
        │   │       channels_out=-1
        │   │   ),
        │   │   accumulate_grad_batches=1,
        │   │   drop_partial_accumulation_on_epoch_end=True,
        │   │   train_budget_steps=inf,
        │   │   train_budget_sec=180.0,
        │   │   train_budget_epochs=inf,
        │   │   lr_schedule=PartialConfig(priml.math.schedules.constant),
        │   │   parallelism=NoParallel.Config(device=None),
        │   │   model_quantization=NoModelQuantization.Config(),
        │   │   activation_memoization=DefaultActivationStorage.Config(),
        │   │   compile=None,
        │   │   ema=NoEMA.Config(),
        │   │   gradient_clip_norm=inf,
        │   │   skip_step_on_nonfinite_grad=False,
        │   │   device_init='eager',
        │   │   dtype_autocast=torch.bfloat16,
        │   │   autocast_cache_enabled=False
        │   ),
        │   dataset=ColdStudies.Config(
        │       preparation=DicomPreparation.Config(
        │           lower=0.02,
        │           upper=0.98,
        │           minimum_span=1e-06,
        │           mean=(0.485, 0.456, 0.406),
        │           std=(0.229, 0.224, 0.225),
        │           base_dir='/opt/scratch',
        │           input_dir=PosixPath('/opt/scratch/datasets/rsna-knee/raw'),
        │           acquisition_dir=PosixPath('/opt/scratch/artifacts/rsna-knee/acquisition'),
        │           size=128,
        │           slices_per_plane=4,
        │           planes=['Sagittal', 'Coronal', 'Axial']
        │       ),
        │       batch_size=2,
        │       validation_fraction=0.25,
        │       split_seed=42,
        │       device=None
        │   ),
        │   metrics_train={},
        │   metrics_eval={'': ProbabilityAuc.Config()},
        │   checkpointer=None,
        │   profiler=None,
        │   phase_timer=PhaseTimer.Config(
        │       enabled=False,
        │       torch_profile=False,
        │       base_dir=PosixPath('/opt/scratch/runs/rsna_knee_training/exp000'),
        │       working_dir=PosixPath('/opt/scratch/runs/rsna_knee_training/exp000/profiling'),
        │       heartbeat_interval_sec=20.0,
        │       fault_dump_interval_sec=40.0,
        │       cuda_events=False
        │   ),
        │   tracker=FileTracker.Config(
        │       base_dir=PosixPath('/opt/scratch/runs/rsna_knee_training/exp000'),
        │       working_dir=PosixPath('/opt/scratch/runs/rsna_knee_training/exp000/metrics.json'),
        │       capture_prefix='eval/'
        │   ),
        │   max_steps=10_000,
        │   max_epochs=inf,
        │   max_time=180.0,
        │   max_time_kind='wall',
        │   max_eval_time=45.0,
        │   eval_stop_on_time_limit=False,
        │   num_steps_eval=-1,
        │   num_steps_log=10,
        │   early_train_log_steps=100,
        │   eval_every_epoch=False,
        │   eval_extras_every_eval=False,
        │   eval_warmup_batches=0,
        │   eval_only=False,
        │   restore_rng_state=True,
        │   seed=None,
        │   model_torch_seed=None,
        │   num_steps_garbage_collect=inf,
        │   mesh_dim_model_seed='pp',
        │   mesh_dim_data_seed='dp',
        │   runtime=SingleProcess.Config(
        │       device='cuda',
        │       deterministic=False,
        │       float32_matmul_precision=None
        │   )
        ),
        budget=WallBudget.Config(
            seconds=300.0,
            reserve_evaluation_sec=45.0,
            clock=<built-in function perf_counter>,
            receipt_path=PosixPath('/opt/scratch/runs/rsna_knee_training/exp000/budget.json')
        ),
        preparation=None,
        export=ColdExport.Config(),
        protected=[PosixPath('/opt/scratch/datasets/rsna-knee/raw')]
    )
