Function all_reduce_norm

yolox/utils/allreduce_norm.py:97–103

All reduce norm statistics in different devices.

(module)

def all_reduce_norm(module):
    """
    All reduce norm statistics in different devices.
    """
    states = get_async_norm_states(module)
    states = all_reduce(states, op="mean")
    module.load_state_dict(states, strict=False)