Add freeze_model_parameters flag to InferenceSettings

Add optional flag to disable requires_grad on all model parameters
during inference. When enabled, only pos and cell retain
requires_grad for autograd-based force/stress computation. The
freeze is applied before torch.compile so the compiled graph can
optimize accordingly. Defaults to False pending further benchmarking.

Author: misko

src/fairchem/core/units/mlip_unit/api/inference.py

@@ -123,6 +123,11 @@ class InferenceSettings:
     # (e.g., eSCNMDBackbone adds stress for all energy tasks by default)
     auto_add_default_untrained_tasks: bool = True
 
+    # When True, set requires_grad=False on all model parameters before
+    # inference. Only pos and cell retain requires_grad for force/stress
+    # computation via autograd.
+    freeze_model_parameters: bool = False
+
     # Maximum number of atoms per system for padding. Required when
     # compile=True for models that use padding (e.g., AllScAIP).
     # All inputs will be padded to this size. Larger values consume more

src/fairchem/core/units/mlip_unit/predict.py

@@ -461,6 +461,12 @@ def _lazy_init(self, data: AtomicData) -> None:
 
         self.move_to_device()
 
+        # Freeze all model parameters before compile — only pos/cell need
+        # requires_grad for force/stress computation via autograd.
+        if self.inference_settings.freeze_model_parameters:
+            for p in self.model.parameters():
+                p.requires_grad_(False)
+
         if self.inference_settings.compile:
             logging.warning(
                 "Model is being compiled this might take a while for the first time"