more flexible assert in test_core.py::test_traceback (#846)

Lightning-AI · Jul 23, 2024 · 6032e4c · 6032e4c
1 parent a16d6ff
commit 6032e4c
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/thunder/tests/test_core.py b/thunder/tests/test_core.py
@@ -1964,8 +1964,10 @@ def f(a):
     with pytest.raises(RuntimeError) as excinfo:
         compiled_f(a)
     assert "on a bool tensor" in str(excinfo.value)
-    assert "torch.neg" in str(excinfo.traceback[-1].statement)
-    assert "thunder.computation" in excinfo.traceback[-1].path
+    # this should actually be in excinfo.traceback[-1] but see
+    # https://github.com/Lightning-AI/lightning-thunder/issues/844
+    assert any(("torch.neg" in str(tb.statement)) for tb in excinfo.traceback)
+    assert any(("thunder.computation" in str(tb.path)) for tb in excinfo.traceback)
 
 
 @instantiate(