emmi_inference.models.loader
============================

.. py:module:: emmi_inference.models.loader


Functions
---------

.. autoapisummary::

   emmi_inference.models.loader.open_checkpoint_file
   emmi_inference.models.loader.load_model


Module Contents
---------------

.. py:function:: open_checkpoint_file(checkpoint_path, device = 'cpu', weights_only = False)

   Open a checkpoint file (regular PyTorch file).

   :param checkpoint_path: Path to a checkpoint file.
   :param device: Device mapping for loading (e.g. `"cpu"`, `"cuda"`).
   :param weights_only: If `True`, open the checkpoint file only for primitive types (dicts, tensors, etc).
                        Defaults to `False`.

   :returns:

             - Dictionary with checkpoint keys and arbitrary values (e.g. state_dict, metadata, etc.).


.. py:function:: load_model(checkpoint, *, model_factory, device = 'cpu', weights_only = False)

   Load a PyTorch model from a checkpoint.

   :param checkpoint: Path to a checkpoint file or a constructor dictionary with relevant fields and values.
   :param model_factory: Factory function to create the model instance
   :param device: PyTorch device for loading (e.g. `"cpu"`, `"cuda"`).
   :param weights_only: passed to torch.load

   :returns: A pytorch model.

   :raises TypeError: in cased of failed state_dict loading into the model.