[BugFix] 0 not into cuda graph to save memory (#5426)

2025-12-24 13:28:13 +08:00 · 2025-12-08 16:47:44 +08:00
parent d1bd40d44c
commit 438c9f785a
3 changed files with 11 additions and 12 deletions
--- a/fastdeploy/distributed/custom_all_reduce/custom_all_reduce.py
+++ b/fastdeploy/distributed/custom_all_reduce/custom_all_reduce.py
@@ -207,6 +207,10 @@ class CustomAllreduce:

    def custom_all_reduce(self, input: paddle.Tensor) -> Optional[paddle.Tensor]:
        """The main allreduce API that provides support for cuda graph."""
+
+        if input.shape[0] == 0:
+            return input
+
        if self.capturing:
            lib = cuda_wrapper.CudaRTLibrary()
            stream = paddle.device.current_stream()