File size: 791 Bytes
742c169 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | import importlib
import torch
from . import env
from .env import TORCH_COMPILE as _ # noqa: F401 (bundles env.py)
def import_class(cls_string):
if not isinstance(cls_string, str):
return cls_string
module, cls = cls_string.rsplit(".", 1)
module = importlib.import_module(module)
return getattr(module, cls)
def compile_wrapper(func):
"""Decorator: compile the function once with torch.compile when TORCH_COMPILE=True."""
_compiled = None
def wrapper(*args, **kwargs):
nonlocal _compiled
if env.TORCH_COMPILE:
if _compiled is None:
_compiled = torch.compile(func, **env.COMPILE_KWARG)
return _compiled(*args, **kwargs)
else:
return func(*args, **kwargs)
return wrapper
|