Query while building one Model in tf2

[Madhusudan_Ghosh]

@Animesh_Sinha

class IndexOfMaxLayer(tf.keras.layers.Layer):
  def __init__(self):
    super(IndexOfMaxLayer, self).__init__()

  def call(self, inputs):
    return tf.math.argmax(inputs)

I want to select one layer from one list (5 layers are already stacked inside this list) for this corresponding index. But this above custom layer returns one keras tensor which does not help to select one layer from that list.