Function initialize_train_state

utils/train_state_utils.py:71–89

(config, model, model_ema , device)

def initialize_train_state(config, model, model_ema , device):
    params = []
    params += model.parameters()
    model_ema.eval()
    logging.warning(f"nnet has {cnt_params(model)} parameters")
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=config.optim.lr, weight_decay=config.optim.wd
    )
  
    train_state = TrainState(
        optimizer=optimizer,
        step=0,
        model=model,
        model_ema=model_ema,
    )
    train_state.ema_update(0)
    if device is not None:
        train_state.to(device)
    return train_state