Function pre_forward_hook

python/oneflow/nn/parallel/distributed.py:158–161

(module, input)

    if broadcast_buffers:

        def pre_forward_hook(module, input):
            with flow.no_grad():
                buffers = list(module.buffers())
                flow._C.comm_broadcast(buffers, inplace=True)

        module.register_forward_pre_hook(pre_forward_hook)