def METHOD_NAME(
    hook_state: Any, bucket: dist.GradBucket, optim_stream_state,
) -> torch.futures.Future[torch.Tensor]:
    # Run original hook
    ddp_weakref = hook_state
    ddp_inst = ddp_weakref()
    reducer, process_group = ddp_inst.reducer, ddp_inst.process_group
    fut = reducer._run_allreduce_hook(bucket)
    optimizer_stream = optim_stream_state.optim_stream
    with torch.cuda.stream(optimizer_stream):
        fut.wait()
        # Apply gradient division since C++ side only allreduces and does
        # not average. TODO: (rohan-varma) the div factor may be different
        # when running with join hook
        bucket.buffer().div_(process_group.size())
        model_params = bucket.parameters()
        grads = bucket.gradients()
        # TODO (rohan-varma): upcast as needed for DDP mixed precision,
        # once optimizer in backward + DDP mixed precision is supported.
        for p, g in zip(model_params, grads):
            if hasattr(p, '_in_backward_optimizers'):
                # Note: need to set grad to the bucket's grad, because
                # running allreduce results in the bucket's grad being
                # reduced, but not grad field.
                if not gradient_is_bucket_view:
                    p.grad = g
                for optim in p._in_backward_optimizers:
                    optim.step()
    # Need to return a Future[Tensor] to obey comm hook API contract.
    ret_fut = torch.futures.Future()
    ret_fut.set_result(bucket.buffer())
    # enqueue a callback to wait for this optimizer stream at the end of
    # backward and set all DDP managed grads to None.
    def wait_for_optim_stream_callback():
        torch.cuda.current_stream().wait_stream(
            optim_stream_state.optim_stream
        )
        # Set DDP managed grads to None
        for param in ddp_inst._get_data_parallel_params(ddp_inst.module):
            if hasattr(param, '_in_backward_optimizers'):
                param.grad = None
        # reset for the next backwards pass
        optim_stream_state.wait_for_optim_stream_enqueued = False
    if not optim_stream_state.wait_for_optim_stream_enqueued:
        Variable._execution_engine.queue_callback(
            wait_for_optim_stream_callback
        )
        # mark that the callback is enqueued
        optim_stream_state.wait_for_optim_stream_enqueued = True
    return ret_fut