Function _execute_by_gpu_num

colossalai/testing/utils.py:202–205

(*args, **kwargs)

    def _wrap_func(f):
        def _execute_by_gpu_num(*args, **kwargs):
            num_avail_gpu = get_accelerator().device_count()
            if num_avail_gpu >= min_gpus:
                f(*args, **kwargs)

        return _execute_by_gpu_num