Hi,
You can use these indices with the other_tensor.gather(dim, indices) function.
Note that you either need to use keepdim=True when you call max or unsqueeze the indices before giving them to gather.