Mike0021/zonos2

1from __future__ import annotations2 3import functools4from contextlib import contextmanager5from typing import TYPE_CHECKING6 7if TYPE_CHECKING:8    import torch9 10 11@contextmanager12def torch_dtype(dtype: torch.dtype):13    import torch  # real import when used14 15    old_dtype = torch.get_default_dtype()16    torch.set_default_dtype(dtype)17    try:18        yield19    finally:20        torch.set_default_dtype(old_dtype)21 22 23def nvtx_annotate(name: str, layer_id_field: str | None = None):24    import torch.cuda.nvtx as nvtx25 26    def decorator(fn):27        @functools.wraps(fn)28        def wrapper(self, *args, **kwargs):29            display_name = name30            if layer_id_field and hasattr(self, layer_id_field):31                display_name = name.format(getattr(self, layer_id_field))32            with nvtx.range(display_name):33                return fn(self, *args, **kwargs)34 35        return wrapper36 37    return decorator38