From 582fb69e2e5e76d8b5e7106d918b0b49275834bf Mon Sep 17 00:00:00 2001
From: Yan Wang <kiya00wy@gmail.com>
Date: Fri, 18 Oct 2024 19:35:39 +0200
Subject: [PATCH] Catch the minimum_cut NetworkXUnbounded exception and give
 clear error message (#1330)

---
 thunder/core/rematerialization.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/thunder/core/rematerialization.py b/thunder/core/rematerialization.py
index 66640ca6a6..f24ee1fad2 100644
--- a/thunder/core/rematerialization.py
+++ b/thunder/core/rematerialization.py
@@ -374,7 +374,13 @@ def add_edges(var):
     g = nx.DiGraph()
     g.add_edges_from(edges)
 
-    _, (reachable, non_reachable) = nx.minimum_cut(g, "source", "sink")
+    try:
+        _, (reachable, non_reachable) = nx.minimum_cut(g, "source", "sink")
+    except Exception:
+        raise RuntimeError(
+            "Failed to compute the min-cut on the graph due to a path with infinite capacity."
+            "Please report this error along with your function and relevant details at: https://github.com/Lightning-AI/lightning-thunder/issues/new"
+        )
 
     cut_edges = set()
     for u, nbrs in ((n, g[n]) for n in reachable):