[BugFix] Fix call to tree.plot in tests

ghstack-source-id: 4a5babbf46294ab6ed4a791e26cfacaf3a41a2e0 Pull Request resolved: #2547
pytorch · Nov 9, 2024 · 09d6866 · 09d6866
1 parent 218d5bf
commit 09d6866
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/test/test_storage_map.py b/test/test_storage_map.py
@@ -381,7 +381,9 @@ def test_forest_build(self):
         r0, *_ = self.dummy_rollouts()
         forest = self._make_forest()
         tree = forest.get_tree(r0[0])
-        tree.plot(make_labels=self.make_labels)
+        for leaf in tree.vertices().values():
+            assert leaf in tree
+        # tree.plot(make_labels=self.make_labels)
 
     def test_forest_vertices(self):
         r0, *_ = self.dummy_rollouts()