code: def compute_memory_usage(
      config_obj, training_set_metadata, model_category
  ) -> int:
      update_config_with_metadata(config_obj, training_set_metadata)
      lm = LudwigModel.create_model(config_obj)
      model_size = lm.get_model_size()  # number of parameters in model
      batch_size = config_obj.trainer.batch_size
      if batch_size == AUTO:
          # Smallest valid batch size that will allow training to complete
          batch_size = MINIMUM_BATCH_SIZE
      memory_usage = (
          model_size
          * (BYTES_PER_WEIGHT + BYTES_OPTIMIZER_PER_WEIGHT)
          * batch_size
      )
      if model_category == TEXT:
          return _get_text_model_memory_usage(
              config_obj.to_dict(), training_set_metadata, memory_usage
          )
      else:
          return memory_usage
  
imports: import copy
  import logging
  import math
  from collections import OrderedDict
  from typing import List
  import psutil
  from ludwig.api import LudwigModel
  from ludwig.backend import initialize_backend
  from ludwig.constants import AUTO, AUTOML_DEFAULT_TEXT_ENCODER, AUTOML_LARGE_TEXT_DATASET, AUTOML_MAX_ROWS_PER_CHECKPOINT, AUTOML_SMALLER_TEXT_ENCODER, AUTOML_SMALLER_TEXT_LENGTH, AUTOML_TEXT_ENCODER_MAX_TOKEN_LEN, HYPEROPT, MINIMUM_BATCH_SIZE, PREPROCESSING, SPACE, TEXT, TRAINER
  from ludwig.data.preprocessing import preprocess_for_training
  from ludwig.features.feature_registries import update_config_with_metadata
  from ludwig.schema.model_config import ModelConfig
  from ludwig.utils.automl.utils import get_model_type
  from ludwig.utils.torch_utils import initialize_pytorch
label_desc: Use image node_id 5 for calling a global function with example usage: compute_memory_usage(config_obj, training_set_metadata, model_category) and returns: _get_text_model_memory_usage, memory_usage