B.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) return the two returns.
is there equivalent method in tensorflow 2.0?