我可以使用张量作为列表索引吗？

我有这个自定义 Keras 层，它可以在列表元素之间进行选择，例如 Dense 层，我希望它直接返回它预测的列表元素。该列表是Keras.layers.Layer. 我有这段代码：

def call(self, inputs, context):   
    pred = tf.argmax(tf.matmul(context, self.kernel))
    return self.layers[pred](inputs)

它抛出一个错误：TypeError: list indices must be integers or slices, not Tensor，这是可以理解的，但我找不到让它工作的方法。“pred” Tensor 没有.numpy属性，尽管我正在急切地运行程序，因为在构建层时会发生这种情况。

我知道可能没有解决方案，如果是这样，请提交有关如何以另一种方式对该层进行编码的想法。