code: def get_configs(model_name="wave_vit_s"):
      args = {
          "wave_vit_s": dict(
              stem_hidden_dim=32,
              embed_dims=[64, 128, 320, 448],
              num_heads=[2, 4, 10, 14],
              mlp_ratios=[8, 8, 4, 4],
              norm_layer=partial(nn.LayerNorm, eps=1e-6),
              depths=[3, 4, 6, 3],
              sr_ratios=[4, 2, 1, 1],
          ),
          "wave_vit_b": dict(
              stem_hidden_dim=64,
              embed_dims=[64, 128, 320, 512],
              num_heads=[2, 4, 10, 16],
              mlp_ratios=[8, 8, 4, 4],
              norm_layer=partial(nn.LayerNorm, eps=1e-6),
              depths=[3, 4, 12, 3],
              sr_ratios=[4, 2, 1, 1],
          ),
          "wave_vit_l": dict(
              stem_hidden_dim=64,
              embed_dims=[96, 192, 384, 512],
              num_heads=[3, 6, 12, 16],
              mlp_ratios=[8, 8, 4, 4],
              norm_layer=partial(nn.LayerNorm, eps=1e-6),
              depths=[3, 6, 18, 3],
              sr_ratios=[4, 2, 1, 1],
          ),
      }
      return args[model_name]
  
imports: import torch
  from torch import nn
  from torch.autograd import Function
  import pywt
  from functools import partial
label_desc: Use image node_id 1 for calling a global function with example usage: get_configs(model_name) and returns: args