Merge branch 'main' of github.com:nimashoghi/MatterTune

Author: lingyukong677

src/mattertune/backbones/orb/model.py

@@ -105,7 +105,7 @@ def requires_disabled_inference_mode(self):
     def _create_output_head(self, prop: props.PropertyConfig, pretrained_model):
         with optional_import_error_message("orb_models"):
             from orb_models.forcefield.forcefield_heads import (
-                EnergyHeadPoolAfter,
+                EnergyHead,
                 ForceHead,
                 StressHead,
                 GraphHead,
@@ -268,7 +268,7 @@ def model_forward_context(self, data, mode: str):
 
 
     @override
-    def model_forward(self, batch, mode: str, using_partition: bool = False):
+    def model_forward(self, batch, mode: str):
         with optional_import_error_message("orb_models"):
             from orb_models.forcefield.forcefield_utils import compute_gradient_forces_and_stress
 
@@ -451,4 +451,4 @@ def create_normalization_context_from_batch(self, batch):
         if compositions is None:
             raise ValueError("No composition found in the batch.")
         compositions = compositions[:, 1:]  # Remove the zeroth element
-        return NormalizationContext(num_atoms=num_atoms, compositions=compositions)
+        return NormalizationContext(num_atoms=num_atoms, compositions=compositions)