core.units.mlip_unit.utils

Copyright (c) Meta Platforms, Inc. and affiliates.

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Functions

load_inference_model(→ tuple[torch.nn.Module, ...)
load_tasks(...)	Load tasks from a checkpoint file.
tf32_context_manager()
update_configs(original_config, new_config)

Module Contents

core.units.mlip_unit.utils.load_inference_model(checkpoint_location: str, overrides: dict | None = None, use_ema: bool = False, return_checkpoint: bool = True) → tuple[torch.nn.Module, fairchem.core.units.mlip_unit.api.inference.MLIPInferenceCheckpoint] | torch.nn.Module

core.units.mlip_unit.utils.load_tasks(checkpoint_location: str) → list[fairchem.core.units.mlip_unit.mlip_unit.Task]

Load tasks from a checkpoint file.

Parameters:

checkpoint_location (str) – Path to the checkpoint file.

Returns:

A list of instantiated Task objects from the checkpoint’s tasks_config.

Return type:

list[Task]

core.units.mlip_unit.utils.tf32_context_manager()

core.units.mlip_unit.utils.update_configs(original_config, new_config)