Fix test more generally (#3165)

Summary:
Pull Request resolved: #3165

https://www.internalfb.com/diff/D77614983
attempted to fix a test, but I still see it showing up in other tests, so this fixes it in general.

Reviewed By: huydhn

Differential Revision: D77758554

fbshipit-source-id: bd390081b68fa650f1cfd6d2a93a1fbf206aaff7

Author: exclamaforte

torchrec/distributed/test_utils/multi_process.py

@@ -90,7 +90,15 @@ def __exit__(self, exc_type, exc_instance, traceback) -> None:
         dist.destroy_process_group(self.pg)
         torch.use_deterministic_algorithms(False)
         if torch.cuda.is_available() and self.disable_cuda_tf_32:
+            # torch/testing/_internal/common_utils.py calls `disable_global_flags()`
+            # workaround RuntimeError: not allowed to set ... after disable_global_flags
+            setattr(  # noqa: B010
+                torch.backends, "__allow_nonbracketed_mutation_flag", True
+            )
             torch.backends.cudnn.allow_tf32 = True
+            setattr(  # noqa: B010
+                torch.backends, "__allow_nonbracketed_mutation_flag", False
+            )
 
 
 class MultiProcessTestBase(unittest.TestCase):

torchrec/distributed/tests/test_pt2_multiprocess.py

@@ -749,11 +749,6 @@ def test_compile_multiprocess(
         model_type, sharding_type, input_type, tovb, compile_backend, config = (
             given_config_tuple
         )
-        # torch/testing/_internal/common_utils.py calls `disable_global_flags()`
-        # workaround RuntimeError: not allowed to set ... after disable_global_flags
-        setattr(  # noqa: B010
-            torch.backends, "__allow_nonbracketed_mutation_flag", True
-        )
         self._run_multi_process_test(
             callable=_test_compile_rank_fn,
             test_model_type=model_type,
@@ -766,9 +761,6 @@ def test_compile_multiprocess(
             config=config,
             torch_compile_backend=compile_backend,
         )
-        setattr(  # noqa: B010
-            torch.backends, "__allow_nonbracketed_mutation_flag", False
-        )
 
     # pyre-ignore
     @unittest.skipIf(