class RNN(tf.keras.Model):
    def __init__(self, embed_dim, hidden_dim, num_layers, vocab_dim, batch_size):
        super(RNN, self).__init__()
        random_init = tf.random_normal_initializer(stddev=0.1)

        self.hidden_dim = hidden_dim
        self.batch_size = batch_size
        self.embedding = tf.Variable(random_init(
            [vocab_dim, embed_dim]), dtype=tf.float32)

        self.gru = tf.keras.layers.GRU(hidden_dim,
                                   return_sequences=True, # return the hidden states
                                   return_state=True)

        self.predict = tf.keras.layers.Dense(vocab_dim)

    def call(self, indices, hidden, training=True):
        indices = self.get_embedding(indices)

        output, state = self.gru(indices, initial_state = hidden)

        # Context vector (batch_size, hidden_size)
        context_vector, attention_weights = attention_layer(state, output)

        # What should I do here?
        preds = self.predict(output)
        return preds, state

    def get_embedding(self, indices):
        return tf.nn.embedding_lookup(self.embedding, indices)

    def initialize_hidden_state(self):
        return tf.zeros((249, self.hidden_dim))