model: R7.Config(
        channels_hidden=384,
        encoder=ConvNeXt.Config(
        │   channels_in=3,
        │   channels_hidden=(96, 192, 384, 768),
        │   blocks_per_stage=(3, 3, 27, 3),
        │   stem=ConvNeXtStem.Config(
        │       channels_in=3,
        │       channels_out=96,
        │       conv=Conv2d.Config(
        │           channels_in=3,
        │           channels_out=96,
        │           kernel_size=4,
        │           stride=4,
        │           padding=0,
        │           dilation=1,
        │           groups=1,
        │           bias=True,
        │           depth_index=(),
        │           device=None,
        │           dtype=None,
        │           init_weight=priml.model.init.kaiming_uniform,
        │           init_bias=torch.nn.init.zeros_
        │       ),
        │       norm=LayerNorm2d.Config(
        │           channels_in=96,
        │           channels_out=96,
        │           eps=1e-06,
        │           elementwise_affine=True,
        │           device=None,
        │           dtype=None
        │       )
        │   ),
        │   stage=[
        │   │   ConvNeXtStage.Config(
        │   │       channels_in=96,
        │   │       channels_out=96,
        │   │       blocks_per_stage=2,
        │   │       downsample=Identity.Config(channels_in=96, channels_out=96),
        │   │       block=[
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=96,
        │   │       │       channels_out=96,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=96,
        │   │       │           channels_out=96,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=96,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=96,
        │   │       │           channels_out=96,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=96,
        │   │       │           channels_out=96,
        │   │       │           channels_hidden=384,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.0
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=96,
        │   │       │       channels_out=96,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=96,
        │   │       │           channels_out=96,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=96,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=96,
        │   │       │           channels_out=96,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=96,
        │   │       │           channels_out=96,
        │   │       │           channels_hidden=384,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.0028571428265422583
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=96,
        │   │       │       channels_out=96,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=96,
        │   │       │           channels_out=96,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=96,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=96,
        │   │       │           channels_out=96,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=96,
        │   │       │           channels_out=96,
        │   │       │           channels_hidden=384,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.0057142856530845165
        │   │       │   )
        │   │       ]
        │   │   ),
        │   │   ConvNeXtStage.Config(
        │   │       channels_in=96,
        │   │       channels_out=192,
        │   │       blocks_per_stage=2,
        │   │       downsample=Sequential.Config(
        │   │           channels_in=96,
        │   │           channels_out=192,
        │   │           elements=[
        │   │               LayerNorm2d.Config(
        │   │                   channels_in=96,
        │   │                   channels_out=96,
        │   │                   eps=1e-06,
        │   │                   elementwise_affine=True,
        │   │                   device=None,
        │   │                   dtype=None
        │   │               ),
        │   │               Conv2d.Config(
        │   │                   channels_in=96,
        │   │                   channels_out=192,
        │   │                   kernel_size=2,
        │   │                   stride=2,
        │   │                   padding=0,
        │   │                   dilation=1,
        │   │                   groups=1,
        │   │                   bias=True,
        │   │                   depth_index=(),
        │   │                   device=None,
        │   │                   dtype=None,
        │   │                   init_weight=priml.model.init.kaiming_uniform,
        │   │                   init_bias=torch.nn.init.zeros_
        │   │               )
        │   │           ],
        │   │           repeat=1,
        │   │           depth_index=()
        │   │       ),
        │   │       block=[
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=192,
        │   │       │       channels_out=192,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=192,
        │   │       │           channels_out=192,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=192,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=192,
        │   │       │           channels_out=192,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=192,
        │   │       │           channels_out=192,
        │   │       │           channels_hidden=768,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.008571428246796131
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=192,
        │   │       │       channels_out=192,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=192,
        │   │       │           channels_out=192,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=192,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=192,
        │   │       │           channels_out=192,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=192,
        │   │       │           channels_out=192,
        │   │       │           channels_hidden=768,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.011428571306169033
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=192,
        │   │       │       channels_out=192,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=192,
        │   │       │           channels_out=192,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=192,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=192,
        │   │       │           channels_out=192,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=192,
        │   │       │           channels_out=192,
        │   │       │           channels_hidden=768,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.014285714365541935
        │   │       │   )
        │   │       ]
        │   │   ),
        │   │   ConvNeXtStage.Config(
        │   │       channels_in=192,
        │   │       channels_out=384,
        │   │       blocks_per_stage=2,
        │   │       downsample=Sequential.Config(
        │   │           channels_in=192,
        │   │           channels_out=384,
        │   │           elements=[
        │   │               LayerNorm2d.Config(
        │   │                   channels_in=192,
        │   │                   channels_out=192,
        │   │                   eps=1e-06,
        │   │                   elementwise_affine=True,
        │   │                   device=None,
        │   │                   dtype=None
        │   │               ),
        │   │               Conv2d.Config(
        │   │                   channels_in=192,
        │   │                   channels_out=384,
        │   │                   kernel_size=2,
        │   │                   stride=2,
        │   │                   padding=0,
        │   │                   dilation=1,
        │   │                   groups=1,
        │   │                   bias=True,
        │   │                   depth_index=(),
        │   │                   device=None,
        │   │                   dtype=None,
        │   │                   init_weight=priml.model.init.kaiming_uniform,
        │   │                   init_bias=torch.nn.init.zeros_
        │   │               )
        │   │           ],
        │   │           repeat=1,
        │   │           depth_index=()
        │   │       ),
        │   │       block=[
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.017142856493592262
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.019999999552965164
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.022857142612338066
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.025714285671710968
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.02857142873108387
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.03142857179045677
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.034285712987184525
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.037142857909202576
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.03999999910593033
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.04285714402794838
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.04571428522467613
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.048571426421403885
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.051428575068712234
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.05428571626543999
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.05714285746216774
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.06000000238418579
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.06285714358091354
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.0657142847776413
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.06857143342494965
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.0714285746216774
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.07428571581840515
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.0771428570151329
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.07999999821186066
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.082857146859169
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.08571428805589676
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.08857142925262451
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=384,
        │   │       │       channels_out=384,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=384,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=384,
        │   │       │           channels_out=384,
        │   │       │           channels_hidden=1_536,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.09142857044935226
        │   │       │   )
        │   │       ]
        │   │   ),
        │   │   ConvNeXtStage.Config(
        │   │       channels_in=384,
        │   │       channels_out=768,
        │   │       blocks_per_stage=2,
        │   │       downsample=Sequential.Config(
        │   │           channels_in=384,
        │   │           channels_out=768,
        │   │           elements=[
        │   │               LayerNorm2d.Config(
        │   │                   channels_in=384,
        │   │                   channels_out=384,
        │   │                   eps=1e-06,
        │   │                   elementwise_affine=True,
        │   │                   device=None,
        │   │                   dtype=None
        │   │               ),
        │   │               Conv2d.Config(
        │   │                   channels_in=384,
        │   │                   channels_out=768,
        │   │                   kernel_size=2,
        │   │                   stride=2,
        │   │                   padding=0,
        │   │                   dilation=1,
        │   │                   groups=1,
        │   │                   bias=True,
        │   │                   depth_index=(),
        │   │                   device=None,
        │   │                   dtype=None,
        │   │                   init_weight=priml.model.init.kaiming_uniform,
        │   │                   init_bias=torch.nn.init.zeros_
        │   │               )
        │   │           ],
        │   │           repeat=1,
        │   │           depth_index=()
        │   │       ),
        │   │       block=[
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=768,
        │   │       │       channels_out=768,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=768,
        │   │       │           channels_out=768,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=768,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=768,
        │   │       │           channels_out=768,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=768,
        │   │       │           channels_out=768,
        │   │       │           channels_hidden=3_072,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.09428571909666061
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=768,
        │   │       │       channels_out=768,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=768,
        │   │       │           channels_out=768,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=768,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=768,
        │   │       │           channels_out=768,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=768,
        │   │       │           channels_out=768,
        │   │       │           channels_hidden=3_072,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.09714286029338837
        │   │       │   ),
        │   │       │   ConvNeXtBlock.Config(
        │   │       │       channels_in=768,
        │   │       │       channels_out=768,
        │   │       │       conv_dw=Conv2d.Config(
        │   │       │           channels_in=768,
        │   │       │           channels_out=768,
        │   │       │           kernel_size=7,
        │   │       │           stride=1,
        │   │       │           padding=3,
        │   │       │           dilation=1,
        │   │       │           groups=768,
        │   │       │           bias=True,
        │   │       │           depth_index=(),
        │   │       │           device=None,
        │   │       │           dtype=None,
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_bias=torch.nn.init.zeros_
        │   │       │       ),
        │   │       │       norm=LayerNorm.Config(
        │   │       │           channels_in=768,
        │   │       │           channels_out=768,
        │   │       │           eps=1e-06,
        │   │       │           elementwise_affine=True,
        │   │       │           device=None,
        │   │       │           dtype=None
        │   │       │       ),
        │   │       │       mlp=SwiGLU.Config(
        │   │       │           channels_in=768,
        │   │       │           channels_out=768,
        │   │       │           channels_hidden=3_072,
        │   │       │           expansion=4.0,
        │   │       │           round_to=1,
        │   │       │           gate=False,
        │   │       │           bias=True,
        │   │       │           split_gate_projection=False,
        │   │       │           norm=None,
        │   │       │           act=priml.model.swiglu.gelu,
        │   │       │           depth_index=(),
        │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │       │           init_weight_out=priml.model.init.kaiming_uniform,
        │   │       │           up_proj=None,
        │   │       │           middle=None,
        │   │       │           down_proj=None,
        │   │       │           shard=None
        │   │       │       ),
        │   │       │       layer_scale=1e-06,
        │   │       │       drop_path=0.10000000149011612
        │   │       │   )
        │   │       ]
        │   │   )
        │   ],
        │   head=Identity.Config(channels_in=768, channels_out=768),
        │   drop_path=0.1,
        │   drop_path_schedule=priml.model.convnext.linear_drop_path
        ),
        proj_in=Linear.Config(
            channels_in=768,
            channels_out=384,
            bias=True,
            depth_index=(),
            device=None,
            dtype=None,
            init_weight=priml.model.init.kaiming_uniform,
            init_bias=torch.nn.init.zeros_,
            linear_fn=<built-in function linear>,
            shard=None
        ),
        norm_in=LayerNorm.Config(
            channels_in=384,
            channels_out=384,
            eps=1e-05,
            elementwise_affine=True,
            device=None,
            dtype=None
        ),
        slot=Embedding.Config(
            channels_in=10,
            channels_out=384,
            padding_idx=None,
            device=None,
            dtype=None,
            lookup_dtype=None,
            shard=None,
            depth_index=(),
            init_weight=torch.nn.init.zeros_
        ),
        view=Embedding.Config(
            channels_in=6,
            channels_out=384,
            padding_idx=None,
            device=None,
            dtype=None,
            lookup_dtype=None,
            shard=None,
            depth_index=(),
            init_weight=torch.nn.init.zeros_
        ),
        geom=Linear.Config(
            channels_in=4,
            channels_out=384,
            bias=False,
            depth_index=(),
            device=None,
            dtype=None,
            init_weight=torch.nn.init.zeros_,
            init_bias=torch.nn.init.zeros_,
            linear_fn=<built-in function linear>,
            shard=None
        ),
        dropout=Dropout.Config(channels_in=384, channels_out=384, p=0.1, inplace=False),
        win_proj=Linear.Config(
            channels_in=384,
            channels_out=384,
            bias=True,
            depth_index=(),
            device=None,
            dtype=None,
            init_weight=priml.model.init.kaiming_uniform,
            init_bias=torch.nn.init.zeros_,
            linear_fn=<built-in function linear>,
            shard=None
        ),
        norm_out=LayerNorm.Config(
            channels_in=384,
            channels_out=384,
            eps=1e-05,
            elementwise_affine=True,
            device=None,
            dtype=None
        ),
        norm_spatial=LayerNorm.Config(
            channels_in=768,
            channels_out=768,
            eps=1e-05,
            elementwise_affine=False,
            device=None,
            dtype=None
        ),
        norm_sp_head=LayerNorm.Config(
            channels_in=768,
            channels_out=768,
            eps=1e-06,
            elementwise_affine=True,
            device=None,
            dtype=None
        ),
        mean=(0.485, 0.456, 0.406),
        std=(0.229, 0.224, 0.225),
        pixel_scale=255.0,
        pad_multiple=32,
        init_weight=functools.partial(priml.model.init.normal, std=0.02),
        init_bias_extent=functools.partial(torch.nn.init.constant_, val=-4.0),
        init_prior_state=functools.partial(torch.nn.init.constant_, val=0.5),
        sp_key=Linear.Config(
            channels_in=768,
            channels_out=128,
            bias=True,
            depth_index=(),
            device=None,
            dtype=None,
            init_weight=priml.model.init.kaiming_uniform,
            init_bias=torch.nn.init.zeros_,
            linear_fn=<built-in function linear>,
            shard=None
        ),
        sp_pos=Linear.Config(
            channels_in=16,
            channels_out=128,
            bias=False,
            depth_index=(),
            device=None,
            dtype=None,
            init_weight=torch.nn.init.zeros_,
            init_bias=torch.nn.init.zeros_,
            linear_fn=<built-in function linear>,
            shard=None
        ),
        frequencies=(1.0, 2.0, 4.0, 8.0),
        state_counts=(6, 6, 5, 5, 6, 6, 6, 6, 5, 6, 6, 8),
        binary_parents=(4, 5, 6, 4, 5, 6, 4, 5, 6, 1, 11, 10, 0),
        extent_classes=(4, 5, 6)
    )
