From 3cd9bb8a4d7bea9c05d8af198095012c302ed1bd Mon Sep 17 00:00:00 2001
From: Ilyes Batatia <48651863+ilyes319@users.noreply.github.com>
Date: Fri, 22 Nov 2024 14:52:46 +0000
Subject: [PATCH] add back interaction first to MACE model

---
 mace/tools/model_script_utils.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/mace/tools/model_script_utils.py b/mace/tools/model_script_utils.py
index 3f49eb41..ad54e2f0 100644
--- a/mace/tools/model_script_utils.py
+++ b/mace/tools/model_script_utils.py
@@ -152,9 +152,7 @@ def _build_model(
             distance_transform=args.distance_transform,
             correlation=args.correlation,
             gate=modules.gate_dict[args.gate],
-            interaction_cls_first=modules.interaction_classes[
-                "RealAgnosticInteractionBlock"
-            ],
+            interaction_cls_first=modules.interaction_classes[args.interaction_first],
             MLP_irreps=o3.Irreps(args.MLP_irreps),
             atomic_inter_scale=args.std,
             atomic_inter_shift=[0.0] * len(heads),