diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index dbcad0056..9e34bbcbc 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -1323,10 +1323,7 @@ class EMA_FSDP: ema.update(model) ema.state_dict() # on rank 0 """ -<<<<<<< HEAD -======= ->>>>>>> 2b2f672c (inference works after changes added) def __init__(self, module, decay: float = 0.999, mode: str = "local_shard"): self.decay = float(decay) self.mode = mode @@ -1420,6 +1417,7 @@ class EMA_FSDP: p.data.copy_(w.to(dtype=p.dtype, device=p.device)) class _ApplyEMACtx: + def __init__(self, ema: "EMA_FSDP", module): self.ema = ema self.module = module