Modify `warnings` in a `with` block to avoid flaky tests (#31893)

* fix

* [test_all] check before merge

---------
Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>

tests/trainer/test_trainer_callback.py

@@ -218,6 +218,7 @@ class TrainerCallbackTest(unittest.TestCase):
         import warnings
 
         # XXX: for now ignore scatter_gather warnings in this test since it's not relevant to what's being tested
+        with warnings.catch_warnings():
             warnings.simplefilter(action="ignore", category=UserWarning)
 
             trainer = self.get_trainer(callbacks=[MyTestTrainerCallback])