Function wrapper

monai/utils/profiling.py:59–65

(*args, **kwargs)

    @wraps(func)
    def wrapper(*args, **kwargs):
        with torch.autograd.profiler.profile(use_cuda=True) as prof:
            result = func(*args, **kwargs)

        print(prof, flush=True)

        return result

    return wrapper