def METHOD_NAME(gpu: int) -> int:
  gpu = int(gpu)
  if 'CUDA_VISIBLE_DEVICES' in os.environ:
    available_devices = os.environ['CUDA_VISIBLE_DEVICES'].split(",")
    if gpu >= len(available_devices):
      raise ValueError(
          f"CUDA_VISIBLE_DEVICES={os.environ['CUDA_VISIBLE_DEVICES']}, arg gpu must be less than {len(available_devices)}"
      )
    return int(available_devices[gpu])
  else:
    return gpu