code: def get_config():
      """Returns the ViT experiment configuration for ImageNet."""
  
      config = ml_collections.ConfigDict()
      config.experiment_name = "vtab_vit_g"
  
      # Dataset (changed in the sweep).
      config.dataset_name = "bit"
      config.data_dtype_str = "float32"
      config.dataset_configs = ml_collections.ConfigDict()
      config.dataset_configs.dataset = ""
      config.dataset_configs.val_split = ()
      config.dataset_configs.train_split = ""
      config.dataset_configs.num_classes = 0
      config.dataset_configs.pp_train = ""
      config.dataset_configs.pp_eval = None
      config.dataset_configs.prefetch_to_device = None
      config.dataset_configs.shuffle_buffer_size = 1_000
      config.batch_size = 512
  
      # Model.
      version, patch = VARIANT.split("/")
      config.model_name = "plainvit"
      config.model = ml_collections.ConfigDict()
      config.model.hidden_size = HIDDEN_SIZES[version]
      config.model.patches = ml_collections.ConfigDict()
      config.model.patches.size = [int(patch), int(patch)]
      config.model.num_heads = NUM_HEADS[version]
      config.model.mlp_dim = MLP_DIMS[version]
      config.model.num_layers = NUM_LAYERS[version]
      config.model.dropout_rate = 0.0
      config.model.classifier = (
          "gap"  # ViT-G was tuned on VTAB using 'gap'.
      )
      config.model.representation_size = None
      config.model.positional_embedding = "learn"
      config.init_head_bias = 0.0
  
      # Pretrained model info.
      config.init_from = ml_collections.ConfigDict()
      config.init_from.checkpoint_format = "big_vision"
      config.init_from.checkpoint_path = CHECKPOINTS[VARIANT]
  
      # Training.
      config.trainer_name = "plainvit_trainer"
      config.loss = "softmax_xent"
      config.l2_decay_factor = None
      config.label_smoothing = None
      config.log_eval_steps = 100
      config.log_summary_steps = 50
      config.num_training_epochs = None  # we use number of steps
      config.num_training_steps = (
          2_500  # ViT-G was tuned on VTAB with 2500 steps.
      )
      config.rng_seed = 42
      sched = ml_collections.ConfigDict()
      sched.re = "(.*)"
      sched.lr_configs = ml_collections.ConfigDict()
      sched.lr_configs.learning_rate_schedule = "compound"
      sched.lr_configs.factors = "constant*linear_warmup*cosine_decay"
      sched.lr_configs.steps_per_cycle = config.num_training_steps
      sched.lr_configs.total_steps = config.num_training_steps
      sched.lr_configs.warmup_steps = (
          200  # ViT-G was tuned on VTAB with 200 warmup steps.
      )
      sched.lr_configs.base_learning_rate = (
          0.01  # ViT-G was tuned on VTAB with 0.01.
      )
  
      # Configure both learning rate schedules.
      config.schedule = ml_collections.ConfigDict({"all": sched})
  
      # *Single* optimizer.
      optim = ml_collections.ConfigDict()
      optim.optax_name = "scenic.momentum_hp"
      # Disable Optax gradient clipping as we handle it ourselves.
      optim.max_grad_norm = 1.0
      optim.weight_decay = 0.0
      config.optimizer = optim
  
      # Logging.
      config.write_summary = True
      config.xprof = True  # Profile using xprof.
      config.checkpoint = True  # Do checkpointing.
      config.checkpoint_steps = 5000
      config.debug_train = False  # Debug mode during training.
      config.debug_eval = False  # Debug mode during eval.
  
      return config
  
imports: import ml_collections
label_desc: Use image node_id 1 for calling a global function with example usage: get_config() and returns: config