Update main.py

main.py

@@ -300,7 +300,10 @@ def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pip
         best_image.save(f"{save_dir}/best_image.png")
         #init_image.save(f"{save_dir}/init_image.png")
 
-        find_and_move_object_to_cpu()
+        # Move the pipe back to CPU
+        if args.model != "flux":
+            # Call the function to find and move object to CPU
+            find_and_move_object_to_cpu()
         
         clear_gpu()