[Buddha7771]

Make EMA track a module's trainable parameters directly, so callers no longer pass tensors on every call.

Motivation

• Every caller had to build the tracked set itself ({n: p for n, p in model.named_parameters() if p.requires_grad})
• update and swap took model.state_dict() on each call, which rebuilds the whole state dict every step only to look tensors up by name
• Names had to match by string, so a mapping from a wrapped module failed only at run time with a key error

Changes

• EMA(module, decay) keeps references to the module's requires_grad parameters and raises if there are none
• update(decay=None) and swap() take no tensors; swap still exchanges in place and swaps back on exit
• state_dict, load_state_dict, and to are unchanged, and keys stay the parameter names, so saved EMA states still load
• Bump the version to 0.2.0 for the breaking change

Migration

# before
ema = EMA({n: p for n, p in model.named_parameters() if p.requires_grad})
ema.update(model.state_dict())
with ema.swap(model.state_dict()):
    ...

# after
ema = EMA(model)
ema.update()
with ema.swap():
    ...

Pass the unwrapped module (not a DDP or torch.compile wrapper) so that names match the module's state_dict.