How to get attention weights in hierarchical model

[saxenarohit]

Model :

sequence_input = Input(shape=(MAX_SENT_LENGTH,), dtype='int32')
words = embedding_layer(sequence_input)
h_words = Bidirectional(GRU(200, return_sequences=True,dropout=0.2,recurrent_dropout=0.2))(words)
sentence = Attention()(h_words)  #with return true
#sentence = Dropout(0.2)(sentence)
sent_encoder = Model(sequence_input, sentence[0])
print(sent_encoder.summary())

document_input = Input(shape=(None, MAX_SENT_LENGTH), dtype='int32')
document_enc = TimeDistributed(sent_encoder)(document_input)
h_sentences = Bidirectional(GRU(100, return_sequences=True))(document_enc)

preds = Dense(7, activation='softmax')(h_sentences)
model = Model(document_input, preds)

Once the model is train, I want to visualise attention weights for a new document.
Since I am using Attention layer(with return attention true) inside time distributed, I am not able to get attention weights

The code I am trying to get attention weights for new input

get_layer_output = K.function([model.layers[0].input, K.learning_phase()],
                                  [model.layers[1].layer.layers[3].output])

layer_output = get_layer_output([new_input, 0])[0]