remove custom export function

metatensor · Oct 24, 2024 · e60fbae · e60fbae
1 parent 47b3798
commit e60fbae
Show file tree

Hide file tree

Showing 8 changed files with 9 additions and 171 deletions.
diff --git a/docs/src/dev-docs/utils/export.rst b/docs/src/dev-docs/utils/export.rst
diff --git a/src/metatrain/experimental/alchemical_model/model.py b/src/metatrain/experimental/alchemical_model/model.py
@@ -7,6 +7,7 @@
 from metatensor.torch.atomistic import (
     MetatensorAtomisticModel,
     ModelCapabilities,
+    ModelMetadata,
     ModelOutput,
     NeighborListOptions,
     System,
@@ -16,7 +17,6 @@
 from ...utils.additive import ZBL
 from ...utils.data.dataset import DatasetInfo
 from ...utils.dtype import dtype_to_str
-from ...utils.export import export
 from .utils import systems_to_torch_alchemical_batch
 
 
@@ -184,7 +184,7 @@ def export(self) -> MetatensorAtomisticModel:
             dtype=dtype_to_str(dtype),
         )
 
-        return export(model=self, model_capabilities=capabilities)
+        return MetatensorAtomisticModel(self.eval(), ModelMetadata(), capabilities)
 
     def set_composition_weights(
         self,

diff --git a/src/metatrain/experimental/gap/model.py b/src/metatrain/experimental/gap/model.py
@@ -14,6 +14,7 @@
 from metatensor.torch.atomistic import (
     MetatensorAtomisticModel,
     ModelCapabilities,
+    ModelMetadata,
     ModelOutput,
     System,
 )
@@ -22,7 +23,6 @@
 from metatrain.utils.data.dataset import DatasetInfo
 
 from ...utils.additive import ZBL, CompositionModel
-from ...utils.export import export
 
 
 class GAP(torch.nn.Module):
@@ -253,7 +253,7 @@ def export(self) -> MetatensorAtomisticModel:
             self._subset_of_regressors.export_torch_script_model()
         )
 
-        return export(model=self, model_capabilities=capabilities)
+        return MetatensorAtomisticModel(self.eval(), ModelMetadata(), capabilities)
 
     def set_composition_weights(
         self,

diff --git a/src/metatrain/experimental/pet/model.py b/src/metatrain/experimental/pet/model.py
@@ -8,6 +8,7 @@
 from metatensor.torch.atomistic import (
     MetatensorAtomisticModel,
     ModelCapabilities,
+    ModelMetadata,
     ModelOutput,
     NeighborListOptions,
     System,
@@ -20,7 +21,6 @@
 
 from ...utils.additive import ZBL
 from ...utils.dtype import dtype_to_str
-from ...utils.export import export
 from .utils import systems_to_batch_dict
 
 
@@ -196,4 +196,4 @@ def export(self) -> MetatensorAtomisticModel:
             supported_devices=["cpu", "cuda"],  # and not __supported_devices__
             dtype=dtype_to_str(dtype),
         )
-        return export(model=self, model_capabilities=capabilities)
+        return MetatensorAtomisticModel(self.eval(), ModelMetadata(), capabilities)
diff --git a/src/metatrain/experimental/soap_bpnn/model.py b/src/metatrain/experimental/soap_bpnn/model.py
@@ -8,6 +8,7 @@
 from metatensor.torch.atomistic import (
     MetatensorAtomisticModel,
     ModelCapabilities,
+    ModelMetadata,
     ModelOutput,
     System,
 )
@@ -18,7 +19,6 @@
 
 from ...utils.additive import ZBL, CompositionModel
 from ...utils.dtype import dtype_to_str
-from ...utils.export import export
 
 
 class Identity(torch.nn.Module):
@@ -335,7 +335,7 @@ def export(self) -> MetatensorAtomisticModel:
             dtype=dtype_to_str(dtype),
         )
 
-        return export(model=self, model_capabilities=capabilities)
+        return MetatensorAtomisticModel(self.eval(), ModelMetadata(), capabilities)
 
     def add_output(self, output_name: str) -> None:
         """Add a new output to the self."""

diff --git a/src/metatrain/utils/export.py b/src/metatrain/utils/export.py
diff --git a/tests/utils/test_export.py b/tests/utils/test_export.py
diff --git a/tox.ini b/tox.ini
@@ -63,7 +63,7 @@ extras =  # architectures used in the package tests
     soap-bpnn
     pet
 allowlist_externals = bash
-#commands_pre = bash {toxinidir}/tests/resources/generate-outputs.sh
+commands_pre = bash {toxinidir}/tests/resources/generate-outputs.sh
 commands =
     pytest \
         --numprocesses=auto \