Make test_create_graph_and_full_backward_hook_cycle more robust to random warnings

soulitzer · soulitzer · commit 3eb04ae4a258 · 2025-05-14T12:54:11.000-04:00
ghstack-source-id: 32e30eb Pull Request resolved: #153407
diff --git a/test/test_autograd.py b/test/test_autograd.py
@@ -9333,12 +9333,11 @@ def forward(self, x):
             tmp.grad_fn.metadata["a"] = test
 
             with set_warn_always_context(True):
-                with warnings.catch_warnings(record=True) as w:
+                with warnings.catch_warnings(record=True) as ws:
                     tmp.exp().sum().backward(create_graph=True)
-                    self.assertTrue(len(w) == 1)
-                    self.assertTrue(
-                        "Using backward() with create_graph=True" in str(w[0].message)
-                    )
+                    filtered_ws = [
+                        w for w in ws if "Using backward() with create_graph=True" in str(w.message)]
+                    self.assertTrue(len(filtered_ws) == 1)
 
             # Remove the backward + create_graph=True cycle
             a.grad = None
