Attention model with seq2seq over sequence
The output of an Attention layer  - the Context - is typically the SUM of the weighted inputs. Each of the input is diminished or magnified by the attention weights based on how relevant it is at that time-step.
So the context will have the same shape as the input. This is typically (batch_size, Encoder_Embedding_dimension). You need to generate this context at every timestep of the decoder. So you will automatically get max_length context's. So the attention layer logic is correct above, you now need to call this in the decoder in a loop.
See &quot;def train_step(inp, targ, enc_hidden):&quot; in the link you have shared where this is done.