Function set_needs_reduction

cdslib/core/utils/multigpu_utils.py:143–144

(self, input, output)

            param.register_hook(allreduce_hook)

    def set_needs_reduction(self, input, output):
        self.needs_reduction = True

    module.register_forward_hook(set_needs_reduction)
    return module