tensorflow2.0中是否有pytorch的最大操作？

Question

overlap-> tensor([[0.0000, 0.0000, 0.0000,  ..., 0.6466, 0.7945, 0.5389]],
       device='cuda:0')

overlap_for_each_prior, object_for_each_prior = overlap.max(dim=0)  # (8732)

this.max(dim=0) 返回两个返回值。
tensorflow 2.0中是否有任何等效方法？

Answer 1

不，你必须使用tf.math.argmax 。 使用它，您还可以获得最大元素：

A = tf.constant([2, 20, 30, 3, 6])
maximum_index = tf.math.argmax(A)
A[maximum_index], maximum_index
# >>> 30, 2