def get_image_version()

in dags/map_reproducibility/utils/common_utils.py [0:0]

• 10 lines of code
• 4 McCabe index (conditional complexity)

def get_image_version(framework: str, model_id: Optional[str] = None):
  if framework == "maxtext":
    return "maxtext_nightly"
  elif framework == "nemo":
    if model_id == "llama3-1-405b":
      return "nemo24.12-A3U"
    else:
      return "nemo24.07-A3U"
  else:
    return None