def METHOD_NAME(self, var_device, var_dtype, apply_state):
    super().METHOD_NAME(var_device, var_dtype, apply_state)
    apply_state[(var_device, var_dtype)]["momentum"] = tf.identity(
        self._get_hyper("momentum", var_dtype)
    )