residual factor

spozdn · May 21, 2024 · 9f6119d · 9f6119d
1 parent 6a3bebb
commit 9f6119d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/default_hypers/default_hypers.yaml b/default_hypers/default_hypers.yaml
@@ -27,6 +27,7 @@ ARCHITECTURAL_HYPERS:
   K_CUT: None # should be float; only used when USE_LONG_RANGE is True
   K_CUT_DELTA: None
   DTYPE: float32 # float32 or float16 or bfloat16
+  RESIDUAL_FACTOR: 0.5
 
 FITTING_SCHEME:  
   INITIAL_LR: 1e-4

diff --git a/src/pet.py b/src/pet.py
@@ -540,6 +540,7 @@ def __init__(self, hypers, transformer_dropout, n_atomic_species):
         self.TARGET_TYPE = hypers.TARGET_TYPE
         self.TARGET_AGGREGATION = hypers.TARGET_AGGREGATION
         self.N_GNN_LAYERS = hypers.N_GNN_LAYERS
+        self.RESIDUAL_FACTOR = hypers.RESIDUAL_FACTOR
 
     def get_predictions(self, batch_dict: Dict[str, torch.Tensor]):
 
@@ -579,7 +580,7 @@ def get_predictions(self, batch_dict: Dict[str, torch.Tensor]):
 
             # batch_dict['input_messages'] = output_messages[neighbors_index, neighbors_pos]
             new_input_messages = output_messages[neighbors_index, neighbors_pos]
-            batch_dict["input_messages"] = 0.5 * (
+            batch_dict["input_messages"] = self.RESIDUAL_FACTOR * (
                 batch_dict["input_messages"] + new_input_messages
             )