Function _max_by_axis

PATH/core/utils.py:743–749

(the_list)

    print(f"[rank {dist.get_rank()}]", *args, **kwargs)

def _max_by_axis(the_list):
    # type: (List[List[int]]) -> List[int]
    maxes = the_list[0]
    for sublist in the_list[1:]:
        for index, item in enumerate(sublist):
            maxes[index] = max(maxes[index], item)
    return maxes


class NestedTensor(object):