Update main.py

main.py

@@ -61,6 +61,7 @@ def setup(args):
     )
     
     torch.cuda.empty_cache()  # Free up cached memory
+    gc.collect()
     
     trainer = LatentNoiseTrainer(
         reward_losses=reward_losses,
@@ -79,6 +80,7 @@ def setup(args):
     )
 
     torch.cuda.empty_cache()  # Free up cached memory
+    gc.collect()
 
     # Create latents
     if args.model == "flux":