Bug fix

Ligo-Biosciences · May 30, 2024 · 5afe517 · 5afe517
1 parent c18aff0
commit 5afe517
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/diffusion/augmentation.py b/src/diffusion/augmentation.py
@@ -22,7 +22,7 @@ def centre_random_augmentation(
     device = atom_positions.x.device
 
     # Center the atoms
-    center = atom_positions.mean(dim=-2, keepdim=True)
+    center = atom_positions.mean(dim=-1, keepdim=True)
     atom_positions = atom_positions - center
 
     # Sample random rotation