Models

Most of the code in metatomic.torch is here to define and export models and to store model metadata. The corresponding classes are documented below:

• Exporting models
• Information about models

We also provide a couple of functions to work with the models:

metatomic.torch.read_model_metadata(path: str) → ModelMetadata

Read metadata of a saved atomistic model.

This function allows to access the metadata of a model without loading it in advance.

Parameters:

path (str) – path to the exported model file

Return type:

ModelMetadata

metatomic.torch.load_atomistic_model(path, extensions_directory=None) → AtomisticModel

Check and then load the metatomic model at the given path.

This function calls metatomic.torch.check_atomistic_model() and metatomic.torch.load_model_extensions() before attempting to load the model.

Parameters:

• path – path to an exported metatensor model

• extensions_directory – path to a directory containing all extensions required by the exported model

Return type:

AtomisticModel

metatomic.torch.check_atomistic_model(path: str)

Check that the file at path contains an exported atomistic model, and that this model can be loaded in the current process.

This function should be called before torch.jit.load() when loading an existing model.

Parameters:

path (str) – path to the exported model file

metatomic.torch.load_model_extensions(path: str, extensions_directory: str | None = None)

Load the TorchScript extensions (and their dependencies) that the model at path uses.

If extensions_directory is provided, we look for the extensions and their dependencies in there first. If this function fails to load some library, it will produce a warning using Torch’s warnings infrastructure. Users can set the METATENSOR_DEBUG_EXTENSIONS_LOADING environment variable to get more informations about a failure in the standard error output.

Parameters:

• path (str) – path to the exported model file

• extensions_directory (str | None) – path to a directory containing the extensions. This directory will typically be created by calling AtomisticModel.export() with collect_extensions=extensions_directory.