I was building a neural network in tensorflow keras and ended up with the following code as a step in the model:

enc = tfkl.Reshape((-1, 20,input_shape[1]))(input_layer)
encoder_output = []
for i in range(enc.shape[1]):
     o = encoder(enc[:, i, ...])
     encoder_output.append(o)
encoder_output = tf.stack(encoder_output, axis=1)

The code WORKS FINE. However it is quite long in the summary and in general not very elegant. I don't know much about functional API and custom layers. But I would like to know if in this case it's best to use a custom layer or build a functional API block.

Thank you