Method update

skills/paper2code/worked/ddpm/src/utils.py:225–231

Update EMA weights after each training step.

(self, model: nn.Module)

    @torch.no_grad()
    def update(self, model: nn.Module):
        """Update EMA weights after each training step."""
        for name, param in model.named_parameters():
            if param.requires_grad and name in self.shadow:
                self.shadow[name].mul_(self.decay).add_(
                    param.data, alpha=1.0 - self.decay
                )

    def apply(self, model: nn.Module):
        """Load EMA weights into model (for evaluation/sampling)."""