使用tf.gather_nd（）选择张量的元素

Question

我有两个张量：

• 形状A [146,33,559]
• 形状B [146,33]

B包含0到559之间的整数，用作indeces。

我所追求的是形状的张量C [146,33]，其中每个元素对应于B给出的指数。

我尝试了tf.gather_nd(A, B)，这给了我错误

InvalidArgumentError (see above for traceback): index innermost dimension length must be <= params rank; saw: 33 vs. 3
 [[Node: GatherNd = GatherNd[Tindices=DT_INT64, Tparams=DT_FLOAT, _device="/job:localhost/replica:0/task:0/cpu:0"](Reshape_34, _recv_contexts_1_0)]]

我还尝试了tf.gather(A, B)，它给了我错误

InvalidArgumentError (see above for traceback): indices[1,0] = 282 is not in [0, 146)
 [[Node: Gather = Gather[Tindices=DT_INT64, Tparams=DT_FLOAT, validate_indices=true, _device="/job:localhost/replica:0/task:0/cpu:0"](Reshape_34, _recv_contexts_1_0)]]

知道如何解决这个问题吗？