add back interaction first to MACE model

ACEsuit · Nov 22, 2024 · 3cd9bb8 · 3cd9bb8
1 parent 38834dc
commit 3cd9bb8
Showing 1 changed file with 1 addition and 3 deletions.
diff --git 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),