Update main.py

main.py

@@ -59,6 +59,9 @@ def setup(args):
     pipe = get_model(
         args.model, dtype, device, args.cache_dir, args.memsave, args.cpu_offloading
     )
+    
+    torch.cuda.empty_cache()  # Free up cached memory
+    
     trainer = LatentNoiseTrainer(
         reward_losses=reward_losses,
         model=pipe,
@@ -75,6 +78,8 @@ def setup(args):
         imageselect=args.imageselect,
     )
 
+    torch.cuda.empty_cache()  # Free up cached memory
+
     # Create latents
     if args.model == "flux":
         # currently only support 512x512 generation
@@ -97,6 +102,7 @@ def setup(args):
             height // pipe.vae_scale_factor,
             width // pipe.vae_scale_factor,
         )
+    
     enable_grad = not args.no_optim
     
     if args.enable_multi_apply:
@@ -111,6 +117,8 @@ def setup(args):
     else:
         multi_apply_fn = None    
 
+    torch.cuda.empty_cache()  # Free up cached memory
+
     return args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings
 
 def execute_task(args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings, progress_callback=None):