diff --git a/unimol/README.md b/unimol/README.md
index 0f9fdb0..6b0b930 100644
--- a/unimol/README.md
+++ b/unimol/README.md
@@ -278,6 +278,12 @@ For QM7, QM8 and QM9, we set `loss_func=finetune_smooth_mae`.
 Molecular conformation generation
 ------------------
 
+**NOTE**: If you would like to reproduce the results from the paper, you can switch to commit 37b0198 or an earlier commit by using the following command:
+```
+git checkout 37b0198cf68a349a854410a06777c2e7dacbce5e
+```
+**Reproduction**
+
 1. Finetune Uni-Mol pretrained model on the training set of the conformation generation task: 
 
 ```bash
diff --git a/unimol/unimol/losses/reg_loss.py b/unimol/unimol/losses/reg_loss.py
index 0f27e77..d567a9c 100644
--- a/unimol/unimol/losses/reg_loss.py
+++ b/unimol/unimol/losses/reg_loss.py
@@ -192,7 +192,7 @@ def reduce_metrics(logging_outputs, split="valid") -> None:
                 .numpy()
                 .mean(axis=1)
             )
-            agg_mae = np.abs(y_pred - y_true).mean(axis=0).mean(axis=1)
+            agg_mae = np.abs(y_pred - y_true).mean()
             metrics.log_scalar(f"{split}_agg_mae", agg_mae, sample_size, round=4)