Update train_diffusion_dpo.py (#6754)

* Update train_diffusion_dpo.py

Address #6702

* Update train_diffusion_dpo_sdxl.py

* Empty-Commit

---------

Co-authored-by: Sayak Paul <[email protected]>

Author: viettmab

examples/research_projects/diffusion_dpo/train_diffusion_dpo.py

@@ -860,7 +860,7 @@ def collate_fn(examples):
                 # Final loss.
                 scale_term = -0.5 * args.beta_dpo
                 inside_term = scale_term * (model_diff - ref_diff)
-                loss = -1 * F.logsigmoid(inside_term.mean())
+                loss = -1 * F.logsigmoid(inside_term).mean()
 
                 implicit_acc = (inside_term > 0).sum().float() / inside_term.size(0)
                 implicit_acc += 0.5 * (inside_term == 0).sum().float() / inside_term.size(0)

examples/research_projects/diffusion_dpo/train_diffusion_dpo_sdxl.py

@@ -975,7 +975,7 @@ def compute_time_ids(original_size, crops_coords_top_left):
                 # Final loss.
                 scale_term = -0.5 * args.beta_dpo
                 inside_term = scale_term * (model_diff - ref_diff)
-                loss = -1 * F.logsigmoid(inside_term.mean())
+                loss = -1 * F.logsigmoid(inside_term).mean()
 
                 implicit_acc = (inside_term > 0).sum().float() / inside_term.size(0)
                 implicit_acc += 0.5 * (inside_term == 0).sum().float() / inside_term.size(0)