| 137 | ) |
| 138 | |
| 139 | def post_forward_hook(module, input, output): |
| 140 | ddp_state_for_reversed_params = module._ddp_state_for_reversed_params |
| 141 | for state in ddp_state_for_reversed_params.values(): |
| 142 | state[0], state[1] = False, False |
| 143 | output = ArgsTree(output).map_leaf( |
| 144 | lambda x: flow._C.select_top_n( |
| 145 | convert_to_tensor_tuple([x, *ddp_state_for_reversed_params.keys()]), |
| 146 | n=1, |
| 147 | )[0] |
| 148 | ) |
| 149 | buffers = list(module.buffers()) |
| 150 | if len(buffers) > 0: |
| 151 | flow._C.stream_touch(buffers) |
| 152 | return output |
| 153 | |
| 154 | module.register_forward_hook(post_forward_hook) |
| 155 | |