fix loss calculation bug

Ligo-Biosciences · Aug 23, 2024 · fcf03e2 · fcf03e2
1 parent 80caef3
commit fcf03e2
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/configs/trainer/deepspeed.yaml b/configs/trainer/deepspeed.yaml
@@ -22,4 +22,4 @@ strategy:
    cpu_checkpointing: False
 
 # Gradient accumulation
-accumulate_grad_batches: 2
+accumulate_grad_batches: 1
diff --git a/src/utils/loss.py b/src/utils/loss.py
@@ -136,7 +136,7 @@ def diffusion_loss(
 
     # Scale by (t**2 + σ**2) / (t + σ)**2
     scaling_factor = torch.add(timesteps ** 2, sd_data ** 2) / (torch.mul(timesteps, sd_data) ** 2 + epsilon)
-    loss_diffusion = scaling_factor * mse
+    loss_diffusion = scaling_factor.squeeze(-1) * mse  # (bs)
 
     # Smooth LDDT Loss
     # if use_smooth_lddt: