def wait_for_gradients():
    """Wait for all gradient reductions to complete and reset bucket state"""
    for handle in bucket_handles:
        if handle is not None:
            handle.wait()

    # Reset state for next iteration
    for i in range(len(bucket_ready_count)):
        bucket_ready_count[i] = 0
        bucket_handles[i] = None
Wait for all gradient reductions to complete and reset bucket state
train_gpt.py