def __init__(self, decay: float, use_num_updates: bool = True):
        """
        decay: The exponential decay.
        use_num_updates: Whether to use number of updates when computing
            averages.
        """
        super().__init__()
        self.decay = decay
        self.use_num_updates = use_num_updates
        self.ema = None

        decay: The exponential decay.
        use_num_updates: Whether to use number of updates when computing
            averages.
        
training/src/callbacks/ema.py