qadence 1.7.3__py3-none-any.whl → 1.7.5__py3-none-any.whl

qadence/ml_tools/config.py

@@ -5,7 +5,7 @@ import os
 from dataclasses import dataclass, field, fields
 from logging import getLogger
 from pathlib import Path
-from typing import Callable, Type
+from typing import Any, Callable, Type
 from uuid import uuid4
 
 from sympy import Basic
@@ -13,6 +13,7 @@ from torch import Tensor
 
 from qadence.blocks.analog import AnalogBlock
 from qadence.blocks.primitive import ParametricBlock
+from qadence.ml_tools.data import OptimizeResult
 from qadence.operations import RX, AnalogRX
 from qadence.parameters import Parameter
 from qadence.types import (
@@ -27,6 +28,84 @@ from qadence.types import (
 
 logger = getLogger(__file__)
 
+CallbackFunction = Callable[[OptimizeResult], None]
+CallbackConditionFunction = Callable[[OptimizeResult], bool]
+
+
+class Callback:
+    """Callback functions are calling in train functions.
+
+    Each callback function should take at least as first input
+    an OptimizeResult instance.
+
+    Attributes:
+        callback (CallbackFunction): Callback function accepting an
+            OptimizeResult as first argument.
+        callback_condition (CallbackConditionFunction | None, optional): Function that
+            conditions the call to callback. Defaults to None.
+        called_every (int, optional): Callback to be called each `called_every` epoch.
+            Defaults to 1.
+            If callback_condition is None, we set
+            callback_condition to returns True when iteration % every == 0.
+        call_before_opt (bool, optional): If true, callback is applied before training.
+            Defaults to False.
+        call_end_epoch (bool, optional): If true, callback is applied during training,
+            after an epoch is performed. Defaults to True.
+        call_after_opt (bool, optional): If true, callback is applied after training.
+            Defaults to False.
+        call_during_eval (bool, optional): If true, callback is applied during evaluation.
+            Defaults to False.
+    """
+
+    def __init__(
+        self,
+        callback: CallbackFunction,
+        callback_condition: CallbackConditionFunction | None = None,
+        called_every: int = 1,
+        call_before_opt: bool = False,
+        call_end_epoch: bool = True,
+        call_after_opt: bool = False,
+        call_during_eval: bool = False,
+    ) -> None:
+        """Initialized Callback.
+
+        Args:
+            callback (CallbackFunction): Callback function accepting an
+                OptimizeResult as ifrst argument.
+            callback_condition (CallbackConditionFunction | None, optional): Function that
+                conditions the call to callback. Defaults to None.
+            called_every (int, optional): Callback to be called each `called_every` epoch.
+                Defaults to 1.
+                If callback_condition is None, we set
+                callback_condition to returns True when iteration % every == 0.
+            call_before_opt (bool, optional): If true, callback is applied before training.
+                Defaults to False.
+            call_end_epoch (bool, optional): If true, callback is applied during training,
+                after an epoch is performed. Defaults to True.
+            call_after_opt (bool, optional): If true, callback is applied after training.
+                Defaults to False.
+            call_during_eval (bool, optional): If true, callback is applied during evaluation.
+                Defaults to False.
+        """
+        self.callback = callback
+        self.call_before_opt = call_before_opt
+        self.call_end_epoch = call_end_epoch
+        self.call_after_opt = call_after_opt
+        self.call_during_eval = call_during_eval
+
+        if called_every <= 0:
+            raise ValueError("Please provide a strictly positive `called_every` argument.")
+        self.called_every = called_every
+
+        if callback_condition is None:
+            self.callback_condition = lambda opt_result: True
+        else:
+            self.callback_condition = callback_condition
+
+    def __call__(self, opt_result: OptimizeResult) -> Any:
+        if opt_result.iteration % self.called_every == 0 and self.callback_condition(opt_result):
+            return self.callback(opt_result)
+
 
 @dataclass
 class TrainConfig:
@@ -45,13 +124,27 @@ class TrainConfig:
     max_iter: int = 10000
     """Number of training iterations."""
     print_every: int = 1000
-    """Print loss/metrics."""
+    """Print loss/metrics.
+
+    Set to 0 to disable
+    """
     write_every: int = 50
-    """Write loss and metrics with the tracking tool."""
+    """Write loss and metrics with the tracking tool.
+
+    Set to 0 to disable
+    """
     checkpoint_every: int = 5000
-    """Write model/optimizer checkpoint."""
+    """Write model/optimizer checkpoint.
+
+    Set to 0 to disable
+    """
     plot_every: int = 5000
-    """Write figures."""
+    """Write figures.
+
+    Set to 0 to disable
+    """
+    callbacks: list[Callback] = field(default_factory=lambda: list())
+    """List of callbacks."""
     log_model: bool = False
     """Logs a serialised version of the model."""
     folder: Path | None = None
@@ -234,13 +327,13 @@ class FeatureMapConfig:
 
     multivariate_strategy: MultivariateStrategy = MultivariateStrategy.PARALLEL
     """
-    The  encoding strategy in case of multi-variate function.
+    The encoding strategy in case of multi-variate function.
 
     Takes qadence.MultivariateStrategy.
     If PARALLEL, the features are encoded in one block of rotation gates
-    with each feature given an equal number of qubits.
-    If SERIES, the features are encoded sequentially, with an ansatz block
-    between. PARALLEL is allowed only for DIGITAL `feature_map_strategy`.
+    with the register being split in sub-registers for each feature.
+    If SERIES, the features are encoded sequentially using the full register for each feature, with
+    an ansatz block between them. PARALLEL is allowed only for DIGITAL `feature_map_strategy`.
     """
 
     feature_map_strategy: Strategy = Strategy.DIGITAL
@@ -261,20 +354,20 @@ class FeatureMapConfig:
     account the domain of the feature-encoding function.
     Defaults to `None` and thus, the feature map is not trainable.
     Note that this is separate from the name of the parameter.
-    The user can provide a single prefix for all features, and they will be appended
+    The user can provide a single prefix for all features, and it will be appended
     by appropriate feature name automatically.
     """
 
     num_repeats: int | dict[str, int] = 0
     """
-    Number of feature map layers repeated in the data reuploadig step.
+    Number of feature map layers repeated in the data reuploading step.
 
-    If all are to be repeated the same number of times, then can give a single
-    `int`. For different number of repeatitions for each feature, provide a dict
+    If all features are to be repeated the same number of times, then can give a single
+    `int`. For different number of repetitions for each feature, provide a dict
     of (str, int) where the key is the name of the variable and the value is the
-    number of repeatitions for that feature.
+    number of repetitions for that feature.
     This amounts to the number of additional reuploads. So if `num_repeats` is N,
-    the data gets uploaded N+1 times. Defaults to no repeatition.
+    the data gets uploaded N+1 times. Defaults to no repetition.
     """
 
     operation: Callable[[Parameter | Basic], AnalogBlock] | Type[RX] | None = None
@@ -300,8 +393,8 @@ class FeatureMapConfig:
         if self.multivariate_strategy == MultivariateStrategy.PARALLEL and self.num_features > 1:
             assert (
                 self.feature_map_strategy == Strategy.DIGITAL
-            ), "For `parallel` encoding of multiple features, the `feature_map_strategy` must be \
-                  of `digital` type."
+            ), "For parallel encoding of multiple features, the `feature_map_strategy` must be \
+                  of `Strategy.DIGITAL`."
 
         if self.operation is None:
             if self.feature_map_strategy == Strategy.DIGITAL:
@@ -314,8 +407,8 @@ class FeatureMapConfig:
                 if isinstance(self.operation, AnalogBlock):
                     logger.warning(
                         "The `operation` is of type `AnalogBlock` but the `feature_map_strategy` is\
-                        `digital`. The `feature_map_strategy` will be modified and given operation\
-                        will be used."
+                        `Strategy.DIGITAL`. The `feature_map_strategy` will be modified and given \
+                        operation will be used."
                     )
 
                     self.feature_map_strategy = Strategy.ANALOG
@@ -324,12 +417,25 @@ class FeatureMapConfig:
                 if isinstance(self.operation, ParametricBlock):
                     logger.warning(
                         "The `operation` is a digital gate but the `feature_map_strategy` is\
-                        `analog`. The `feature_map_strategy` will be modified and given operation\
-                        will be used."
+                        `Strategy.ANALOG`. The `feature_map_strategy` will be modified and given\
+                        operation will be used."
                     )
 
                     self.feature_map_strategy = Strategy.DIGITAL
 
+            elif self.feature_map_strategy == Strategy.RYDBERG:
+                if self.operation is not None:
+                    logger.warning(
+                        f"feature_map_strategy is `Strategy.RYDBERG` which does not take any\
+                        operation. But an operation {self.operation} is provided. The \
+                        `feature_map_strategy` will be modified and given operation will be used."
+                    )
+
+                    if isinstance(self.operation, AnalogBlock):
+                        self.feature_map_strategy = Strategy.ANALOG
+                    else:
+                        self.feature_map_strategy = Strategy.DIGITAL
+
         if self.inputs is not None:
             assert (
                 len(self.inputs) == self.num_features
@@ -388,16 +494,16 @@ class AnsatzConfig:
     ansatz_type: AnsatzType = AnsatzType.HEA
     """What type of ansatz.
 
-    HEA for Hardware Efficient Ansatz.
-    IIA for Identity intialized Ansatz.
+    `AnsatzType.HEA` for Hardware Efficient Ansatz.
+    `AnsatzType.IIA` for Identity intialized Ansatz.
     """
 
     ansatz_strategy: Strategy = Strategy.DIGITAL
     """Ansatz strategy.
 
-    DIGITAL for fully digital ansatz. Required if `ansatz_type` is `iia`.
-    SDAQC for analog entangling block.
-    RYDBERG for fully rydberg hea ansatz.
+    `Strategy.DIGITAL` for fully digital ansatz. Required if `ansatz_type` is `AnsatzType.IIA`.
+    `Strategy.SDAQC` for analog entangling block.
+    `Strategy.RYDBERG` for fully rydberg hea ansatz.
     """
 
     strategy_args: dict = field(default_factory=dict)
@@ -406,7 +512,7 @@ class AnsatzConfig:
 
     Details about each below.
 
-    For DIGITAL strategy, accepts the following:
+    For `Strategy.DIGITAL` strategy, accepts the following:
         periodic (bool): if the qubits should be linked periodically.
             periodic=False is not supported in emu-c.
         operations (list): list of operations to cycle through in the
@@ -417,7 +523,7 @@ class AnsatzConfig:
             will have variational parameters on the rotation angles.
             Defaults to CNOT
 
-    For SDAQC strategy, accepts the following:
+    For `Strategy.SDAQC` strategy, accepts the following:
         operations (list): list of operations to cycle through in the
             digital single-qubit rotations of each layer.
             Defaults to  [RX, RY, RX] for hea and [RX, RY] for iia.
@@ -425,7 +531,7 @@ class AnsatzConfig:
             analog entangling layer. Time parameter is considered variational.
             Defaults to NN interaction.
 
-    For RYDBERG strategy, accepts the following:
+    For `Strategy.RYDBERG` strategy, accepts the following:
         addressable_detuning: whether to turn on the trainable semi-local addressing pattern
             on the detuning (n_i terms in the Hamiltonian).
             Defaults to True.

qadence/ml_tools/constructors.py

@@ -45,7 +45,7 @@ from .models import QNN
 def _create_support_arrays(
     num_qubits: int,
     num_features: int,
-    multivariate_strategy: str,
+    multivariate_strategy: MultivariateStrategy,
 ) -> list[tuple[int, ...]]:
     """
     Create the support arrays for the digital feature map.
@@ -53,8 +53,8 @@ def _create_support_arrays(
     Args:
         num_qubits (int): The number of qubits.
         num_features (int): The number of features.
-        multivariate_strategy (str): The multivariate encoding strategy.
-            Either 'series' or 'parallel'.
+        multivariate_strategy (MultivariateStrategy): The multivariate encoding strategy.
+            Either 'MultivariateStrategy.SERIES' or 'MultivariateStrategy.PARALLEL'.
 
     Returns:
         list[tuple[int, ...]]: The list of support arrays. ith element of the list is the support
@@ -62,23 +62,25 @@ def _create_support_arrays(
 
     Raises:
         ValueError: If the number of features is greater than the number of qubits
-            with parallel encoding. Not possible to encode these features in parallel.
-        ValueError: If the multivariate strategy is not 'series' or 'parallel'.
+            and the strategy is `MultivariateStrategy.PARALLEL` not possible to assign a support
+            array to each feature.
+        ValueError: If the multivariate strategy is not `MultivariateStrategy.SERIES` or
+            `MultivariateStrategy.PARALLEL`.
     """
-    if multivariate_strategy == "series":
+    if multivariate_strategy == MultivariateStrategy.SERIES:
         return [tuple(range(num_qubits)) for i in range(num_features)]
-    elif multivariate_strategy == "parallel":
+    elif multivariate_strategy == MultivariateStrategy.PARALLEL:
         if num_features <= num_qubits:
             return [tuple(x.tolist()) for x in np.array_split(np.arange(num_qubits), num_features)]
         else:
             raise ValueError(
                 f"Number of features {num_features} must be less than or equal to the number of \
-                qubits {num_qubits}. if the features are to be encoded is parallely."
+                qubits {num_qubits} if the features are to be encoded parallely."
             )
     else:
         raise ValueError(
-            f"Invalid encoding strategy {multivariate_strategy} provided. Only 'series' or \
-                'parallel' are allowed."
+            f"Invalid encoding strategy {multivariate_strategy} provided. Only \
+                `MultivariateStrategy.SERIES` or `MultivariateStrategy.PARALLEL` are allowed."
         )
 
 
@@ -212,7 +214,8 @@ def _create_digital_fm(
         list[AbstractBlock]: The list of digital feature map blocks.
 
     Raises:
-        ValueError: If the encoding strategy is invalid. Only 'series' or 'parallel' are allowed.
+        ValueError: If the encoding strategy is invalid. Only `MultivariateStrategy.SERIES` or
+            `MultivariateStrategy.PARALLEL` are allowed.
     """
     if config.multivariate_strategy == MultivariateStrategy.SERIES:
         fm_blocks = _encode_features_series_digital(register, config)
@@ -220,8 +223,8 @@ def _create_digital_fm(
         fm_blocks = _encode_features_parallel_digital(register, config)
     else:
         raise ValueError(
-            f"Invalid encoding strategy {config.multivariate_strategy} provided. Only 'series' or \
-                'parallel' are allowed."
+            f"Invalid encoding strategy {config.multivariate_strategy} provided. Only\
+            `MultivariateStrategy.SERIES` or `MultivariateStrategy.PARALLEL` are allowed."
         )
 
     return fm_blocks
@@ -348,7 +351,8 @@ def create_fm_blocks(
         list[AbstractBlock]: A list of feature map blocks.
 
     Raises:
-        ValueError: If the feature map strategy is not 'digital', 'analog' or 'rydberg'.
+        ValueError: If the feature map strategy is not `Strategy.DIGITAL`, `Strategy.ANALOG` or
+            `Strategy.RYDBERG`.
     """
     if config.feature_map_strategy == Strategy.DIGITAL:
         return _create_digital_fm(register=register, config=config)
@@ -359,7 +363,7 @@ def create_fm_blocks(
     else:
         raise NotImplementedError(
             f"Feature map not implemented for strategy {config.feature_map_strategy}. \
-            Only 'digital', 'analog' or 'rydberg' allowed."
+            Only `Strategy.DIGITAL`, `Strategy.ANALOG` or `Strategy.RYDBERG` allowed."
         )
 
 
@@ -463,7 +467,8 @@ def _create_iia(
         AbstractBlock: The Identity Initialized Ansatz.
 
     Raises:
-        ValueError: If the ansatz strategy is not supported. Only 'digital' and 'sdaqc' are allowed.
+        ValueError: If the ansatz strategy is not supported. Only `Strategy.DIGITAL` and
+            `Strategy.SDAQC` are allowed.
     """
     if config.ansatz_strategy == Strategy.DIGITAL:
         return _create_iia_digital(num_qubits=num_qubits, config=config)
@@ -471,8 +476,8 @@ def _create_iia(
         return _create_iia_sdaqc(num_qubits=num_qubits, config=config)
     else:
         raise ValueError(
-            f"Invalid ansatz strategy {config.ansatz_strategy} provided. Only 'digital', 'sdaqc', \
-            allowed for IIA."
+            f"Invalid ansatz strategy {config.ansatz_strategy} provided. Only `Strategy.DIGITAL` \
+                and `Strategy.SDAQC` allowed for IIA."
         )
 
 
@@ -487,7 +492,7 @@ def _create_hea_digital(num_qubits: int, config: AnsatzConfig) -> AbstractBlock:
     Returns:
         AbstractBlock: The Digital Hardware Efficient Ansatz.
     """
-    operations = config.strategy_args.get("rotations", [RX, RY, RX])
+    operations = config.strategy_args.get("operations", [RX, RY, RX])
     entangler = config.strategy_args.get("entangler", CNOT)
     periodic = config.strategy_args.get("periodic", False)
 
@@ -512,7 +517,7 @@ def _create_hea_sdaqc(num_qubits: int, config: AnsatzConfig) -> AbstractBlock:
     Returns:
         AbstractBlock: The SDAQC Hardware Efficient Ansatz.
     """
-    operations = config.strategy_args.get("rotations", [RX, RY, RX])
+    operations = config.strategy_args.get("operations", [RX, RY, RX])
     entangler = config.strategy_args.get(
         "entangler", hamiltonian_factory(num_qubits, interaction=Interaction.NN)
     )
@@ -556,7 +561,7 @@ def _create_hea_rydberg(
     )
 
 
-def _create_hea_ansatz(
+def _create_hea(
     register: int | Register,
     config: AnsatzConfig,
 ) -> AbstractBlock:
@@ -571,7 +576,8 @@ def _create_hea_ansatz(
         AbstractBlock: The hardware efficient ansatz block.
 
     Raises:
-        ValueError: If the ansatz strategy is not 'digital', 'sdaqc', or 'rydberg'.
+        ValueError: If the ansatz strategy is not `Strategy.DIGITAL`, `Strategy.SDAQC`, or
+            `Strategy.RYDBERG`.
     """
     num_qubits = register if isinstance(register, int) else register.n_qubits
 
@@ -583,8 +589,8 @@ def _create_hea_ansatz(
         return _create_hea_rydberg(register=register, config=config)
     else:
         raise ValueError(
-            f"Invalid ansatz strategy {config.ansatz_strategy} provided. Only 'digital', 'sdaqc', \
-                and 'rydberg' allowed"
+            f"Invalid ansatz strategy {config.ansatz_strategy} provided. Only `Strategy.DIGITAL`, \
+                `Strategy.SDAQC`, and `Strategy.RYDBERG` allowed"
         )
 
 
@@ -610,11 +616,11 @@ def create_ansatz(
     if config.ansatz_type == AnsatzType.IIA:
         return _create_iia(num_qubits=num_qubits, config=config)
     elif config.ansatz_type == AnsatzType.HEA:
-        return _create_hea_ansatz(register=register, config=config)
+        return _create_hea(register=register, config=config)
     else:
         raise NotImplementedError(
-            f"Ansatz of type {config.ansatz_type} not implemented yet. Only 'hea' and\
-                'iia' available."
+            f"Ansatz of type {config.ansatz_type} not implemented yet. Only `AnsatzType.HEA` and\
+                `AnsatzType.IIA` available."
         )
 
 
@@ -651,7 +657,7 @@ def load_observable_transformations(config: ObservableConfig) -> tuple[Parameter
         config (ObservableConfig): Observable configuration.
 
     Returns:
-        tuple[float, float]: The observable shifting and scaling factors.
+        tuple[Parameter, Parameter]: The observable shifting and scaling factors.
     """
     shift = config.shift
     scale = config.scale
@@ -665,9 +671,9 @@ def load_observable_transformations(config: ObservableConfig) -> tuple[Parameter
 
 
 ObservableTransformMap = {
-    ObservableTransform.RANGE: lambda detuning, scale, shift: (shift, shift - scale)
-    if detuning is N
-    else (0.5 * (shift - scale), 0.5 * (scale + shift)),
+    ObservableTransform.RANGE: lambda detuning, scale, shift: (
+        (shift, shift - scale) if detuning is N else (0.5 * (shift - scale), 0.5 * (scale + shift))
+    ),
     ObservableTransform.SCALE: lambda _, scale, shift: (scale, shift),
 }
 
@@ -718,8 +724,8 @@ def create_observable(
     """
     if transformation_type == ObservableTransform.RANGE:
         scale, shift = ObservableTransformMap[transformation_type](detuning, scale, shift)  # type: ignore[index]
-    shifting_term = shift * _global_identity(register)  # type: ignore[operator]
-    detuning_hamiltonian = scale * hamiltonian_factory(  # type: ignore[operator]
+    shifting_term: AbstractBlock = shift * _global_identity(register)  # type: ignore[operator]
+    detuning_hamiltonian: AbstractBlock = scale * hamiltonian_factory(  # type: ignore[operator]
         register=register,
         detuning=detuning,
     )

qadence/ml_tools/data.py

@@ -1,15 +1,41 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from functools import singledispatch
 from itertools import cycle
 from typing import Any, Iterator
 
+from nevergrad.optimization.base import Optimizer as NGOptimizer
 from torch import Tensor
 from torch import device as torch_device
+from torch.nn import Module
+from torch.optim import Optimizer
 from torch.utils.data import DataLoader, IterableDataset, TensorDataset
 
 
+@dataclass
+class OptimizeResult:
+    """OptimizeResult stores many optimization intermediate values.
+
+    We store at a current iteration,
+    the model, optimizer, loss values, metrics. An extra dict
+    can be used for saving other information to be used for callbacks.
+    """
+
+    iteration: int
+    """Current iteration number."""
+    model: Module
+    """Model at iteration."""
+    optimizer: Optimizer | NGOptimizer
+    """Optimizer at iteration."""
+    loss: Tensor | float | None = None
+    """Loss value."""
+    metrics: dict = field(default_factory=lambda: dict())
+    """Metrics that can be saved during training."""
+    extra: dict = field(default_factory=lambda: dict())
+    """Extra dict for saving anything else to be used in callbacks."""
+
+
 @dataclass
 class DictDataLoader:
     """This class only holds a dictionary of `DataLoader`s and samples from them."""

qadence/ml_tools/optimize_step.py

@@ -29,10 +29,11 @@ def optimize_step(
         xs (dict | list | torch.Tensor | None): the input data. If None it means
             that the given model does not require any input data
         device (torch.device): A target device to run computation on.
+        dtype (torch.dtype): Data type for xs conversion.
 
     Returns:
-        tuple: tuple containing the model, the optimizer, a dictionary with
-            the collected metrics and the compute value loss
+        tuple: tuple containing the computed loss value, and a dictionary with
+            the collected metrics.
     """
 
     loss, metrics = None, {}

qadence/ml_tools/printing.py

@@ -4,7 +4,6 @@ from logging import getLogger
 from typing import Any, Callable, Union
 
 from matplotlib.figure import Figure
-from mlflow.models import infer_signature
 from torch import Tensor
 from torch.nn import Module
 from torch.utils.data import DataLoader
@@ -82,6 +81,7 @@ def plot_mlflow(
 def log_model_mlflow(
     writer: Any, model: Module, dataloader: DataLoader | DictDataLoader | None
 ) -> None:
+    signature = None
     if dataloader is not None:
         xs: InputData
         xs, *_ = next(iter(dataloader))
@@ -94,9 +94,17 @@ def log_model_mlflow(
                 xs[key] = val.numpy()
             for key, val in preds.items():
                 preds[key] = val.detach.numpy()
-        signature = infer_signature(xs, preds)
-    else:
-        signature = None
+
+        try:
+            from mlflow.models import infer_signature
+
+            signature = infer_signature(xs, preds)
+        except ImportError:
+            logger.warning(
+                "An MLFlow specific function has been called but MLFlow failed to import."
+                "Please install MLFlow or adjust your code."
+            )
+
     writer.pytorch.log_model(model, artifact_path="model", signature=signature)
 
 

qadence/ml_tools/saveload.py

@@ -72,7 +72,8 @@ def write_checkpoint(
     device = None
     try:
         # We extract the device from the pyqtorch native circuit
-        device = str(model.device).split(":")[0]  # in case of using several CUDA devices
+        device = model.device if isinstance(QuantumModel, QNN) else next(model.parameters()).device
+        device = str(device).split(":")[0]  # in case of using several CUDA devices
     except Exception as e:
         msg = (
             f"Unable to identify in which device the QuantumModel is stored due to {e}."
@@ -132,7 +133,7 @@ def load_model(
     try:
         iteration, model_dict = torch.load(folder / model_ckpt_name, *args, **kwargs)
         if isinstance(model, (QuantumModel, QNN)):
-            model._from_dict(model_dict, as_torch=True)
+            model.load_params_from_dict(model_dict)
         elif isinstance(model, Module):
             model.load_state_dict(model_dict, strict=True)
         # Load model to a specific gpu device if specified