mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
27 lines
526 B
Python
27 lines
526 B
Python
from contextlib import contextmanager
|
|
from functools import wraps
|
|
|
|
from tensorrt_llm import profiler
|
|
|
|
|
|
def profile(tag):
|
|
|
|
@contextmanager
|
|
def profile_range(tag):
|
|
profiler.start(tag)
|
|
yield
|
|
profiler.stop(tag)
|
|
|
|
def inner_decorator(func):
|
|
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
"""A wrapper function"""
|
|
with profile_range(tag):
|
|
func(*args, **kwargs)
|
|
profiler.summary()
|
|
|
|
return wrapper
|
|
|
|
return inner_decorator
|