split model setup and task execution

main.py

@@ -14,7 +14,8 @@ from rewards import get_reward_losses
 from training import LatentNoiseTrainer, get_optimizer
 
 
-def main(args, progress_callback=None):
+def setup(args):
+    #args = parse_args()
     seed_everything(args.seed)
     bf.makedirs(f"{args.save_dir}/logs/{args.task}")
     # Set up logging and name settings
@@ -92,6 +93,10 @@ def main(args, progress_callback=None):
         )
     enable_grad = not args.no_optim
 
+    return args, trainer, device, dtype, shape, enable_grad, settings
+
+def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, progress_callback=None):
+    #args = parse_args()
     if args.task == "single":
         init_latents = torch.randn(shape, device=device, dtype=dtype)
         latents = torch.nn.Parameter(init_latents, requires_grad=enable_grad)
@@ -269,7 +274,12 @@ def main(args, progress_callback=None):
     # log total rewards
     logging.info(f"Mean initial rewards: {total_init_rewards}")
     logging.info(f"Mean best rewards: {total_best_rewards}")
-    
-if __name__ == "__main__":
+
+def main():
     args = parse_args()
-    main(args)
+    args, trainer, device, dtype, shape, enable_grad, settings = setup(args)
+    execute_task(args, trainer, device, dtype, shape, enable_grad, settings)
+
+
+if __name__ == "__main__":
+    main()