Code refactoring

Signed-off-by: Daniel Korzekwa <[email protected]>

Author: danielkorzekwa

tests/experimental/torch/_compress/nas/plugins/test_nas_convert.py

@@ -48,11 +48,9 @@ def test_nas_convert(project_root_path: Path, tmp_path: Path):
 def _test_nas_convert_multiprocess_job(
     project_root_path: Path, tmp_path: Path, rank: int, size: int
 ):
-    runtime = NativeDdpRuntime(
+    with NativeDdpRuntime(
         dtype=torch.bfloat16, torch_distributed_timeout=datetime.timedelta(10)
-    )
-
-    with runtime as runtime:
+    ) as runtime:
         converted_model, puzzle_dir = run_nas_convert(project_root_path, tmp_path, rank, runtime)
 
         #

tests/experimental/torch/_compress/nas/plugins/test_nas_search.py

@@ -40,11 +40,9 @@ def test_nas_search(project_root_path: Path, tmp_path: Path):
 def _test_nas_search_multiprocess_job(
     project_root_path: Path, tmp_path: Path, rank: int, size: int
 ):
-    runtime = NativeDdpRuntime(
+    with NativeDdpRuntime(
         dtype=torch.bfloat16, torch_distributed_timeout=datetime.timedelta(10)
-    )
-
-    with runtime as runtime:
+    ) as runtime:
         converted_model, puzzle_dir = run_nas_convert(project_root_path, tmp_path, rank, runtime)
 
         #