Unverified Commit 13cc7f53 authored by Lucas Wilkinson's avatar Lucas Wilkinson Committed by GitHub
Browse files

[BugFix] Fix DBO hang (#25625)


Signed-off-by: default avatarLucas Wilkinson <lwilkins@redhat.com>
parent 916bd920
......@@ -366,7 +366,8 @@ class UBatchWrapper:
cudagraph_runtime_mode=CUDAGraphMode.NONE)
with self.sm_control:
return self._capture_ubatches(ubatch_metadata, self.model)
elif num_tokens in self.cudagraphs:
elif num_tokens in self.cudagraphs \
and cudagraph_runtime_mode is CUDAGraphMode.FULL:
cudagraph_metadata = self.cudagraphs[num_tokens]
cudagraph_metadata.cudagraph.replay()
return cudagraph_metadata.outputs
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment