[Bugfix][Compile] Fix gc.collect/empty_cache patch arity in CUDAGraphWrapper (#41235)

Signed-off-by: Roi Koren <roik@nvidia.com>
2026-04-29 21:59:18 +00:00
parent ccfb620c62
commit c2fb013312
1 changed files with 7 additions and 2 deletions
@@ -290,9 +290,14 @@ class CUDAGraphWrapper:
                    # across layers will make the cudagraph capture very slow.
                    # therefore, we only run gc for the first graph,
                    # and disable gc for the rest of the graphs.
-                    stack.enter_context(patch("gc.collect", lambda: None))
                    stack.enter_context(
-                        patch("torch.accelerator.empty_cache", lambda: None)
+                        patch("gc.collect", lambda *args, **kwargs: None)
+                    )
+                    stack.enter_context(
+                        patch(
+                            "torch.accelerator.empty_cache",
+                            lambda *args, **kwargs: None,
+                        )
                    )

                if self.graph_pool is not None: