From 5ed0e78a5643f74f3a749c625d3ea453763b3301 Mon Sep 17 00:00:00 2001
From: zhiyili1230 <zhiyi@orbitalmaterials.com>
Date: Fri, 4 Oct 2024 11:38:09 +0000
Subject: [PATCH] use v2 default

---
 orb_models/forcefield/gns.py             | 4 ++--
 orb_models/forcefield/graph_regressor.py | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/orb_models/forcefield/gns.py b/orb_models/forcefield/gns.py
index 88d9eed..e117794 100644
--- a/orb_models/forcefield/gns.py
+++ b/orb_models/forcefield/gns.py
@@ -219,8 +219,8 @@ def __init__(
         num_edge_out: int,
         num_mlp_layers: int,
         mlp_hidden_dim: int,
-        attention_gate: Literal["sigmoid", "softmax"] = "sigmoid",
-        distance_cutoff: bool = True,
+        attention_gate: Literal["sigmoid", "softmax"] = "softmax",
+        distance_cutoff: bool = False,
         polynomial_order: Optional[int] = 4,
         cutoff_rmax: Optional[float] = 6.0,
     ):
diff --git a/orb_models/forcefield/graph_regressor.py b/orb_models/forcefield/graph_regressor.py
index 42878c5..2679c8a 100644
--- a/orb_models/forcefield/graph_regressor.py
+++ b/orb_models/forcefield/graph_regressor.py
@@ -117,7 +117,7 @@ def __init__(
         mlp_hidden_dim: int,
         target: Union[str, PropertyDefinition],
         dropout: Optional[float] = None,
-        remove_mean: bool = True,
+        remove_mean: bool = False,
     ):
         """Initializes the NodeHead MLP.