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)

    All reduce norm statistics in different devices.
    
yolox/utils/allreduce_norm.py