Followup fix for Custom AR V2 in non NVL scenarios (#24742)

Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>
This commit is contained in:
Brayden Zhong
2026-05-10 16:57:46 -07:00
committed by GitHub
co-authored by b8zhong
parent b202778aa5
commit 8acb0270fd
@@ -93,12 +93,14 @@ class CustomAllReduceV2:
@contextmanager
def capture(self):
if self.disabled:
yield
return
try:
self.obj.set_cuda_graph_capture(True)
yield
finally:
self.obj.set_cuda_graph_capture(False)
if not self.disabled:
# cannot call when graph is capturing
assert (
torch.cuda.is_current_stream_capturing() == False