update tests for LaunchResults

Author: apoorvkh

tests/test_ci.py

@@ -37,8 +37,7 @@ def dist_func() -> torch.Tensor:
         backend="gloo",  # log_dir="./test_logs"
     )
 
-    results = next(iter(r.values()))
-    assert torch.all(results[0] == results[1])
+    assert torch.all(r.value(0) == r.value(1))
 
 
 def test_logging() -> None:

tests/test_func.py

@@ -13,7 +13,7 @@ def test_launch() -> None:
         workers_per_host="slurm",
     )
 
-    result_values = [v for host_results in result.values() for v in host_results.values()]
+    result_values = result.all(by='rank')
 
     t = True
     for i in range(len(result_values)):