Method __init__

utils/train_state_utils.py:22–26

(self, optimizer, step, model=None, model_ema=None)

class TrainState(object):
    def __init__(self, optimizer, step, model=None, model_ema=None):
        self.optimizer = optimizer
        self.step = step
        self.model = model
        self.model_ema = model_ema

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