def determine_context(device_ids: List[int],
                      use_cpu: bool,
                      disable_device_locking: bool,
                      lock_dir: str,
                      exit_stack: ExitStack) -> List[mx.Context]:
    """
    Determine the MXNet context to run on (CPU or GPU).

    :param device_ids: List of device as defined from the CLI.
    :param use_cpu: Whether to use the CPU instead of GPU(s).
    :param disable_device_locking: Disable Sockeye's device locking feature.
    :param lock_dir: Directory to place device lock files in.
    :param exit_stack: An ExitStack from contextlib.
    :return: A list with the context(s) to run on.
    """
    if use_cpu:
        context = [mx.cpu()]
    else:
        num_gpus = get_num_gpus()
        check_condition(num_gpus >= 1,
                        "No GPUs found, consider running on the CPU with --use-cpu ")
        if disable_device_locking:
            context = expand_requested_device_ids(device_ids)
        else:
            context = exit_stack.enter_context(acquire_gpus(device_ids, lock_dir=lock_dir))
        context = [mx.gpu(gpu_id) for gpu_id in context]
    return context
Determine the MXNet context to run on (CPU or GPU).

    :param device_ids: List of device as defined from the CLI.
    :param use_cpu: Whether to use the CPU instead of GPU(s).
    :param disable_device_locking: Disable Sockeye's device locking feature.
    :param lock_dir: Directory to place device lock files in.
    :param exit_stack: An ExitStack from contextlib.
    :return: A list with the context(s) to run on.