def init_state(self, x):
    """
    Initialize t, m, and u
    """
    optim_state = {}
    optim_state["t"] = 0.
    optim_state["m"] = [tf.zeros_like(v) for v in x]
    optim_state["u"] = [tf.zeros_like(v) for v in x]
    return optim_state