From 5afe517d675af583bbbefe5d65533d3f3a81536e Mon Sep 17 00:00:00 2001
From: ardagoreci <62720042+ardagoreci@users.noreply.github.com>
Date: Thu, 30 May 2024 01:17:45 +0100
Subject: [PATCH] Bug fix

---
 src/diffusion/augmentation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/diffusion/augmentation.py b/src/diffusion/augmentation.py
index 21138eb..69f0a31 100644
--- 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