Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
fix pynccl worker_fn_with_cudagraph
Browse files Browse the repository at this point in the history
cennn committed Jan 5, 2025
1 parent 55fe076 commit a11a6de
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion tests/distributed/test_pynccl.py
Original file line number Diff line number Diff line change
@@ -137,7 +137,7 @@ def worker_fn_with_cudagraph():
# run something in the default stream to initialize torch engine
a = torch.ones((4, 4), device=f'cuda:{pynccl_comm.rank}')
torch.cuda.synchronize()
with torch.cuda.graph(graph, stream=torch.cuda.current_stream()), \
with torch.cuda.graph(graph), \
pynccl_comm.change_state(enable=True):
a_out = pynccl_comm.all_reduce(a)
torch.cuda.synchronize()

0 comments on commit a11a6de

Please sign in to comment.