code: def get_train_step(train_step_name: str, **train_step_kwargs):
      """
      Lookup the train_step_name in the train step registry and return.
      If the train step is not implemented, asserts will be thrown and workflow will exit.
      """
      assert (
          train_step_name in TRAIN_STEP_REGISTRY
      ), "Unknown train step"
      return TRAIN_STEP_REGISTRY[train_step_name](**train_step_kwargs)
  
imports: from pathlib import Path
  from classy_vision.generic.registry_utils import import_all_modules
label_desc: Use image node_id 2 for calling a global function with example usage: get_train_step(train_step_name) and returns: TRAIN_STEP_REGISTRY