def _streaming_external_state(self, inputs, state1, state2):
    # first dimension is batch size
    if inputs.shape[0] != self.inference_batch_size:
      raise ValueError(
          'inputs.shape[0]:%d must be = self.inference_batch_size:%d' %
          (inputs.shape[0], self.inference_batch_size))

    # receive inputs: [batch, 1, feature]
    # convert it for lstm cell to inputs1: [batch, feature]
    inputs1 = tf.keras.backend.squeeze(inputs, axis=1)
    output, states = self.lstm_cell(inputs1, [state1, state2])

    # output [batch, 1, feature]
    output = tf.keras.backend.expand_dims(output, axis=1)
    return output, states[0], states[1]