added weighted alignment for the loss

Ligo-Biosciences · Aug 13, 2024 · 7bd3f73 · 7bd3f73
1 parent aa60790
commit 7bd3f73
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/utils/loss.py b/src/utils/loss.py
@@ -109,9 +109,8 @@ def diffusion_loss(
     pred_atoms = Vec3Array.from_array(pred_atoms)
     gt_atoms = Vec3Array.from_array(gt_atoms)
 
-    # Align the gt_atoms to pred_atoms  TODO: add the alignment back in
-    aligned_gt_atoms = gt_atoms
-    # weighted_rigid_align(x=gt_atoms, x_gt=pred_atoms, weights=weights, mask=mask))
+    # Align the gt_atoms to pred_atoms
+    aligned_gt_atoms = weighted_rigid_align(x=gt_atoms, x_gt=pred_atoms, weights=weights, mask=mask)
 
     # MSE loss
     mse = mean_squared_error(pred_atoms, aligned_gt_atoms, weights, mask)