Function enable_stratified

k_diffusion/utils.py:284–294

A context manager that enables stratified sampling.

(group=0, groups=1, disable=False)

@contextmanager
def enable_stratified(group=0, groups=1, disable=False):
    """A context manager that enables stratified sampling."""
    try:
        stratified_settings.disable = disable
        stratified_settings.group = group
        stratified_settings.groups = groups
        yield
    finally:
        del stratified_settings.disable
        del stratified_settings.group
        del stratified_settings.groups


@contextmanager