code: def get_model(model_name, batch_size, num_classes):
      """Gets model_name from TF Model Garden."""
      if model_name == "resnet":
          return load_resnet(batch_size, num_classes)
      elif model_name == "efficientnet":
          return load_efficientnet(num_classes)
      elif model_name in efficientnet_model.MODEL_CONFIGS:
          return load_efficientnet(num_classes, model_name)
      raise TypeError(f"Unknown model_name argument: {model_name}")
  
imports: import os
  import pickle
  from typing import Any, Dict, Optional
  import uuid
  from ..None import machine_config
  from ..None import run
  import tensorflow
  import tensorflow_datasets
  from official.legacy.image_classification.efficientnet import efficientnet_model
  from official.legacy.image_classification.resnet import resnet_model
label_desc: Use image node_id 6 for calling a global function with example usage: get_model(model_name, batch_size, num_classes) and returns: load_resnet, load_efficientnet, load_efficientnet