Function fsdp_ema_setup

train/fsdp_utils.py:550–555

(ema_model, fsdp_config, ignored_modules=[])

def fsdp_ema_setup(ema_model, fsdp_config, ignored_modules=[]):
    for param in ema_model.parameters():
        param.requires_grad = False

    ema_model = fsdp_wrapper(ema_model, fsdp_config, ignored_modules=ignored_modules)
    return ema_model


@torch.no_grad()