[Feature] Add EMA checkpoint

### What is the feature?

I wanted to evaluate ema checkpoint that get from EMAHook , but it saved only in base checkpoint and inside `state_dict` , I wanted to evaluate it and got several bugs! please save it like before in `epoch_*_ema.pth` part. 

### Any other context?

_No response_