def set_gpu(gpu_mon, args):
    """
    Sets the GPU visibility based on the passed arguments. Takes an already
    initialized GPUMonitor object. Sets GPUs according to args.force_GPU, if
    specified, otherwise sets first args.num_GPUs free GPUs on the system.

    Stops the GPUMonitor process once GPUs have been set
    If gpu_mon is None, this function does nothing

    Args:
        gpu_mon: An initialized GPUMonitor object or None
        args: argparse arguments

    Returns: The number of GPUs that was actually set (different from
    args.num_GPUs if args.force_GPU is set to more than 1 GPU)
    """
    num_GPUs = args.num_GPUs
    if gpu_mon is not None:
        if not args.force_GPU:
            gpu_mon.await_and_set_free_GPU(N=num_GPUs, sleep_seconds=120)
        else:
            gpu_mon.set_GPUs = args.force_GPU
            num_GPUs = len(args.force_GPU.split(","))
        gpu_mon.stop()
    return num_GPUs