Method ema_update

utils/train_state_utils.py:28–30

(self, rate=0.9999)

        self.model_ema = model_ema

    def ema_update(self, rate=0.9999):
        if self.model_ema is not None:
            ema(self.model_ema, self.model, rate)

    def save(self, path):
        os.makedirs(path, exist_ok=True)