Pulls out save_validation_info boolean logic

Signed-off-by: Rafael Vasquez <rafvasq21@gmail.com>

Author: rafvasq

aiu_fms_testing_utils/scripts/drive_paged_programs.py

@@ -349,9 +349,6 @@ def get_distributed_kwargs(
         aiu_dist_setup(dist.get_rank(), dist.get_world_size())
         distributed_kwargs["distributed_strategy"] = "tp"
         distributed_kwargs["group"] = dist.group.WORLD
-        save_validation_info_outputs = save_validation_info_outputs and (
-            dist.get_rank() == 0
-        )
 
     return distributed_kwargs
 
@@ -812,6 +809,9 @@ def main():
     distributed_kwargs = get_distributed_kwargs(
         args.distributed, args.dist_timeout, args.save_validation_info_outputs
     )
+    args.save_validation_info_outputs = args.save_validation_info_outputs and (
+        dist.get_rank() == 0
+    )
 
     model = load_model(
         device_type="cpu",