qadence 1.7.2__tar.gz → 1.7.3__tar.gz

{qadence-1.7.2 → qadence-1.7.3}/.gitignore

@@ -75,3 +75,7 @@ events.out.tfevents.*
 *.dvi
 
 *.gv
+
+# mlflow
+mlruns/
+mlartifacts/

{qadence-1.7.2 → qadence-1.7.3}/PKG-INFO

@@ -1,8 +1,8 @@
 Metadata-Version: 2.3
 Name: qadence
-Version: 1.7.2
+Version: 1.7.3
 Summary: Pasqal interface for circuit-based quantum computing SDKs
-Author-email: Aleksander Wennersteen <aleksander.wennersteen@pasqal.com>, Gert-Jan Both <gert-jan.both@pasqal.com>, Niklas Heim <niklas.heim@pasqal.com>, Mario Dagrada <mario.dagrada@pasqal.com>, Vincent Elfving <vincent.elfving@pasqal.com>, Dominik Seitz <dominik.seitz@pasqal.com>, Roland Guichard <roland.guichard@pasqal.com>, "Joao P. Moutinho" <joao.moutinho@pasqal.com>, Vytautas Abramavicius <vytautas.abramavicius@pasqal.com>, Gergana Velikova <gergana.velikova@pasqal.com>, Eduardo Maschio <eduardo.maschio@pasqal.com>, Smit Chaudhary <smit.chaudhary@pasqal.com>, Ignacio Fernández Graña <ignacio.fernandez-grana@pasqal.com>, Charles Moussa <charles.moussa@pasqal.com>
+Author-email: Aleksander Wennersteen <aleksander.wennersteen@pasqal.com>, Gert-Jan Both <gert-jan.both@pasqal.com>, Niklas Heim <niklas.heim@pasqal.com>, Mario Dagrada <mario.dagrada@pasqal.com>, Vincent Elfving <vincent.elfving@pasqal.com>, Dominik Seitz <dominik.seitz@pasqal.com>, Roland Guichard <roland.guichard@pasqal.com>, "Joao P. Moutinho" <joao.moutinho@pasqal.com>, Vytautas Abramavicius <vytautas.abramavicius@pasqal.com>, Gergana Velikova <gergana.velikova@pasqal.com>, Eduardo Maschio <eduardo.maschio@pasqal.com>, Smit Chaudhary <smit.chaudhary@pasqal.com>, Ignacio Fernández Graña <ignacio.fernandez-grana@pasqal.com>, Charles Moussa <charles.moussa@pasqal.com>, Giorgio Tosti Balducci <giorgio.tosti-balducci@pasqal.com>
 License: Apache 2.0
 License-File: LICENSE
 Classifier: License :: OSI Approved :: Apache Software License
@@ -22,7 +22,7 @@ Requires-Dist: matplotlib
 Requires-Dist: nevergrad
 Requires-Dist: numpy
 Requires-Dist: openfermion
-Requires-Dist: pyqtorch==1.2.5
+Requires-Dist: pyqtorch==1.3.2
 Requires-Dist: pyyaml
 Requires-Dist: rich
 Requires-Dist: scipy
@@ -33,6 +33,7 @@ Requires-Dist: torch
 Provides-Extra: all
 Requires-Dist: braket; extra == 'all'
 Requires-Dist: libs; extra == 'all'
+Requires-Dist: mlflow; extra == 'all'
 Requires-Dist: protocols; extra == 'all'
 Requires-Dist: pulser; extra == 'all'
 Requires-Dist: visualization; extra == 'all'
@@ -51,6 +52,8 @@ Requires-Dist: optax; extra == 'horqrux'
 Requires-Dist: sympy2jax; extra == 'horqrux'
 Provides-Extra: libs
 Requires-Dist: qadence-libs; extra == 'libs'
+Provides-Extra: mlflow
+Requires-Dist: mlflow; extra == 'mlflow'
 Provides-Extra: protocols
 Requires-Dist: qadence-protocols; extra == 'protocols'
 Provides-Extra: pulser

{qadence-1.7.2 → qadence-1.7.3}/pyproject.toml

@@ -21,10 +21,11 @@ authors = [
     { name = "Smit Chaudhary", email = "smit.chaudhary@pasqal.com" },
     { name = "Ignacio Fernández Graña", email = "ignacio.fernandez-grana@pasqal.com" },
     { name = "Charles Moussa", email = "charles.moussa@pasqal.com" },
+    { name = "Giorgio Tosti Balducci", email = "giorgio.tosti-balducci@pasqal.com" },
 ]
 requires-python = ">=3.9"
 license = { text = "Apache 2.0" }
-version = "1.7.2"
+version = "1.7.3"
 classifiers = [
     "License :: OSI Approved :: Apache Software License",
     "Programming Language :: Python",
@@ -48,7 +49,7 @@ dependencies = [
     "jsonschema",
     "nevergrad",
     "scipy",
-    "pyqtorch==1.2.5",
+    "pyqtorch==1.3.2",
     "pyyaml",
     "matplotlib",
     "Arpeggio==2.0.2",
@@ -83,8 +84,8 @@ horqrux = [
 protocols = ["qadence-protocols"]
 libs = ["qadence-libs"]
 dlprof = ["nvidia-pyindex", "nvidia-dlprof[pytorch]"]
-all = ["pulser", "braket", "visualization", "protocols", "libs"]
-
+mlflow = ["mlflow"]
+all = ["pulser", "braket", "visualization", "protocols", "libs", "mlflow"]
 
 [tool.hatch.envs.default]
 dependencies = [
@@ -102,7 +103,7 @@ dependencies = [
     "ruff",
     "pydocstringformatter",
 ]
-features = ["pulser", "braket", "visualization", "horqrux"]
+features = ["pulser", "braket", "visualization", "horqrux", "mlflow"]
 
 [tool.hatch.envs.default.scripts]
 test = "pytest -n auto --cov-report lcov --cov-config=pyproject.toml --cov=qadence --cov=tests --ignore=./tests/test_examples.py {args}"
@@ -139,7 +140,7 @@ dependencies = [
     "markdown-exec",
     "mike",
 ]
-features = ["pulser", "braket", "horqrux", "visualization"]
+features = ["pulser", "braket", "horqrux", "visualization", "mlflow"]
 
 [tool.hatch.envs.docs.scripts]
 build = "mkdocs build --clean --strict"

{qadence-1.7.2 → qadence-1.7.3}/qadence/backends/pyqtorch/convert_ops.py

@@ -8,6 +8,7 @@ import pyqtorch as pyq
 import sympy
 import torch
 from pyqtorch.apply import apply_operator
+from pyqtorch.embed import Embedding
 from pyqtorch.matrices import _dagger
 from pyqtorch.time_dependent.sesolve import sesolve
 from pyqtorch.utils import is_diag
@@ -24,7 +25,7 @@ from torch import (
 )
 from torch import device as torch_device
 from torch import dtype as torch_dtype
-from torch.nn import Module
+from torch.nn import Module, ParameterDict
 
 from qadence.backends.utils import (
     finitediff,
@@ -200,7 +201,12 @@ class PyQHamiltonianEvolution(Module):
         elif isinstance(block.generator, Tensor):
             m = block.generator.to(dtype=cdouble)
             hmat = block_to_tensor(
-                MatrixBlock(m, qubit_support=block.qubit_support),
+                MatrixBlock(
+                    m,
+                    qubit_support=block.qubit_support,
+                    check_unitary=False,
+                    check_hermitian=True,
+                ),
                 qubit_support=self.qubit_support,
                 use_full_support=False,
             )
@@ -313,7 +319,8 @@ class PyQHamiltonianEvolution(Module):
     def forward(
         self,
         state: Tensor,
-        values: dict[str, Tensor],
+        values: dict[str, Tensor] | ParameterDict = dict(),
+        embedding: Embedding | None = None,
     ) -> Tensor:
         if getattr(self.block.generator, "is_time_dependent", False):  # type: ignore [union-attr]
 

{qadence-1.7.2 → qadence-1.7.3}/qadence/blocks/matrix.py

@@ -60,7 +60,13 @@ class MatrixBlock(PrimitiveBlock):
     name = "MatrixBlock"
     matrix: torch.Tensor
 
-    def __init__(self, matrix: torch.Tensor | np.ndarray, qubit_support: tuple[int, ...]) -> None:
+    def __init__(
+        self,
+        matrix: torch.Tensor | np.ndarray,
+        qubit_support: tuple[int, ...],
+        check_unitary: bool = True,
+        check_hermitian: bool = False,
+    ) -> None:
         if isinstance(matrix, np.ndarray):
             matrix = torch.tensor(matrix)
         if matrix.ndim == 3 and matrix.size(0) == 1:
@@ -69,10 +75,12 @@ class MatrixBlock(PrimitiveBlock):
             raise TypeError("Please provide a 2D matrix.")
         if not self.is_square(matrix):
             raise TypeError("Please provide a square matrix.")
-        if not self.is_hermitian(matrix):
-            logger.warning("Provided matrix is not hermitian.")
-        if not self.is_unitary(matrix):
-            logger.warning("Provided matrix is not unitary.")
+        if check_hermitian:
+            if not self.is_hermitian(matrix):
+                logger.warning("Provided matrix is not hermitian.")
+        if check_unitary:
+            if not self.is_unitary(matrix):
+                logger.warning("Provided matrix is not unitary.")
         self.matrix = matrix.clone()
         super().__init__(qubit_support)
 

{qadence-1.7.2 → qadence-1.7.3}/qadence/engines/torch/differentiable_expectation.py

@@ -151,8 +151,9 @@ class DifferentiableExpectation:
         return (
             AdjointExpectation.apply(
                 self.circuit.native,
-                self.observable[0].native,  # Currently, adjoint only supports a single observable.
                 self.state,
+                self.observable[0].native,  # Currently, adjoint only supports a single observable.
+                None,
                 self.param_values.keys(),
                 *self.param_values.values(),
             )

{qadence-1.7.2 → qadence-1.7.3}/qadence/ml_tools/config.py

@@ -5,15 +5,25 @@ import os
 from dataclasses import dataclass, field, fields
 from logging import getLogger
 from pathlib import Path
-from typing import Callable, Optional, Type
+from typing import Callable, Type
+from uuid import uuid4
 
 from sympy import Basic
+from torch import Tensor
 
 from qadence.blocks.analog import AnalogBlock
 from qadence.blocks.primitive import ParametricBlock
 from qadence.operations import RX, AnalogRX
 from qadence.parameters import Parameter
-from qadence.types import AnsatzType, BasisSet, MultivariateStrategy, ReuploadScaling, Strategy
+from qadence.types import (
+    AnsatzType,
+    BasisSet,
+    ExperimentTrackingTool,
+    LoggablePlotFunction,
+    MultivariateStrategy,
+    ReuploadScaling,
+    Strategy,
+)
 
 logger = getLogger(__file__)
 
@@ -37,10 +47,14 @@ class TrainConfig:
     print_every: int = 1000
     """Print loss/metrics."""
     write_every: int = 50
-    """Write tensorboard logs."""
+    """Write loss and metrics with the tracking tool."""
     checkpoint_every: int = 5000
     """Write model/optimizer checkpoint."""
-    folder: Optional[Path] = None
+    plot_every: int = 5000
+    """Write figures."""
+    log_model: bool = False
+    """Logs a serialised version of the model."""
+    folder: Path | None = None
     """Checkpoint/tensorboard logs folder."""
     create_subfolder_per_run: bool = False
     """Checkpoint/tensorboard logs stored in subfolder with name `<timestamp>_<PID>`.
@@ -59,14 +73,38 @@ class TrainConfig:
 
     validation loss across previous iterations.
     """
-    validation_criterion: Optional[Callable] = None
+    validation_criterion: Callable | None = None
     """A boolean function which evaluates a given validation metric is satisfied."""
-    trainstop_criterion: Optional[Callable] = None
+    trainstop_criterion: Callable | None = None
     """A boolean function which evaluates a given training stopping metric is satisfied."""
     batch_size: int = 1
     """The batch_size to use when passing a list/tuple of torch.Tensors."""
     verbose: bool = True
     """Whether or not to print out metrics values during training."""
+    tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+    """The tracking tool of choice."""
+    hyperparams: dict = field(default_factory=dict)
+    """Hyperparameters to track."""
+    plotting_functions: tuple[LoggablePlotFunction, ...] = field(default_factory=tuple)  # type: ignore
+    """Functions for in-train plotting."""
+
+    # tensorboard only allows for certain types as hyperparameters
+    _tb_allowed_hyperparams_types: tuple = field(
+        default=(int, float, str, bool, Tensor), init=False, repr=False
+    )
+
+    def _filter_tb_hyperparams(self) -> None:
+        keys_to_remove = [
+            key
+            for key, value in self.hyperparams.items()
+            if not isinstance(value, TrainConfig._tb_allowed_hyperparams_types)
+        ]
+        if keys_to_remove:
+            logger.warning(
+                f"Tensorboard cannot log the following hyperparameters: {keys_to_remove}."
+            )
+            for key in keys_to_remove:
+                self.hyperparams.pop(key)
 
     def __post_init__(self) -> None:
         if self.folder:
@@ -81,6 +119,64 @@ class TrainConfig:
             self.trainstop_criterion = lambda x: x <= self.max_iter
         if self.validation_criterion is None:
             self.validation_criterion = lambda *x: False
+        if self.hyperparams and self.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
+            self._filter_tb_hyperparams()
+        if self.tracking_tool == ExperimentTrackingTool.MLFLOW:
+            self._mlflow_config = MLFlowConfig()
+        if self.plotting_functions and self.tracking_tool != ExperimentTrackingTool.MLFLOW:
+            logger.warning("In-training plots are only available with mlflow tracking.")
+        if not self.plotting_functions and self.tracking_tool == ExperimentTrackingTool.MLFLOW:
+            logger.warning("Tracking with mlflow, but no plotting functions provided.")
+
+    @property
+    def mlflow_config(self) -> MLFlowConfig:
+        if self.tracking_tool == ExperimentTrackingTool.MLFLOW:
+            return self._mlflow_config
+        else:
+            raise AttributeError(
+                "mlflow_config is available only for with the mlflow tracking tool."
+            )
+
+
+class MLFlowConfig:
+    """
+    Configuration for mlflow tracking.
+
+    Example:
+
+        export MLFLOW_TRACKING_URI=tracking_uri
+        export MLFLOW_EXPERIMENT=experiment_name
+        export MLFLOW_RUN_NAME=run_name
+    """
+
+    def __init__(self) -> None:
+        import mlflow
+
+        self.tracking_uri: str = os.getenv("MLFLOW_TRACKING_URI", "")
+        """The URI of the mlflow tracking server.
+
+        An empty string, or a local file path, prefixed with file:/.
+        Data is stored locally at the provided file (or ./mlruns if empty).
+        """
+
+        self.experiment_name: str = os.getenv("MLFLOW_EXPERIMENT", str(uuid4()))
+        """The name of the experiment.
+
+        If None or empty, a new experiment is created with a random UUID.
+        """
+
+        self.run_name: str = os.getenv("MLFLOW_RUN_NAME", str(uuid4()))
+        """The name of the run."""
+
+        mlflow.set_tracking_uri(self.tracking_uri)
+
+        # activate existing or create experiment
+        exp_filter_string = f"name = '{self.experiment_name}'"
+        if not mlflow.search_experiments(filter_string=exp_filter_string):
+            mlflow.create_experiment(name=self.experiment_name)
+
+        self.experiment = mlflow.set_experiment(self.experiment_name)
+        self.run = mlflow.start_run(run_name=self.run_name, nested=False)
 
 
 @dataclass

qadence-1.7.3/qadence/ml_tools/printing.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+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
+from torch.utils.tensorboard import SummaryWriter
+
+from qadence.ml_tools.data import DictDataLoader
+from qadence.types import ExperimentTrackingTool
+
+logger = getLogger(__name__)
+
+PlottingFunction = Callable[[Module, int], tuple[str, Figure]]
+InputData = Union[Tensor, dict[str, Tensor]]
+
+
+def print_metrics(loss: float | None, metrics: dict, iteration: int) -> None:
+    msg = " ".join(
+        [f"Iteration {iteration: >7} | Loss: {loss:.7f} -"]
+        + [f"{k}: {v.item():.7f}" for k, v in metrics.items()]
+    )
+    print(msg)
+
+
+def write_tensorboard(
+    writer: SummaryWriter, loss: float = None, metrics: dict = {}, iteration: int = 0
+) -> None:
+    if loss is not None:
+        writer.add_scalar("loss", loss, iteration)
+    for key, arg in metrics.items():
+        writer.add_scalar(key, arg, iteration)
+
+
+def log_hyperparams_tensorboard(writer: SummaryWriter, hyperparams: dict, metrics: dict) -> None:
+    writer.add_hparams(hyperparams, metrics)
+
+
+def plot_tensorboard(
+    writer: SummaryWriter,
+    model: Module,
+    iteration: int,
+    plotting_functions: tuple[PlottingFunction],
+) -> None:
+    for pf in plotting_functions:
+        descr, fig = pf(model, iteration)
+        writer.add_figure(descr, fig, global_step=iteration)
+
+
+def log_model_tensorboard(
+    writer: SummaryWriter,
+    model: Module,
+    dataloader: Union[None, DataLoader, DictDataLoader],
+) -> None:
+    logger.warning("Model logging is not supported by tensorboard. No model will be logged.")
+
+
+def write_mlflow(writer: Any, loss: float | None, metrics: dict, iteration: int) -> None:
+    writer.log_metrics({"loss": float(loss)}, step=iteration)  # type: ignore
+    writer.log_metrics(metrics, step=iteration)  # logs the single metrics
+
+
+def log_hyperparams_mlflow(writer: Any, hyperparams: dict, metrics: dict) -> None:
+    writer.log_params(hyperparams)  # type: ignore
+
+
+def plot_mlflow(
+    writer: Any,
+    model: Module,
+    iteration: int,
+    plotting_functions: tuple[PlottingFunction],
+) -> None:
+    for pf in plotting_functions:
+        descr, fig = pf(model, iteration)
+        writer.log_figure(fig, descr)
+
+
+def log_model_mlflow(
+    writer: Any, model: Module, dataloader: DataLoader | DictDataLoader | None
+) -> None:
+    if dataloader is not None:
+        xs: InputData
+        xs, *_ = next(iter(dataloader))
+        preds = model(xs)
+        if isinstance(xs, Tensor):
+            xs = xs.numpy()
+            preds = preds.detach().numpy()
+        elif isinstance(xs, dict):
+            for key, val in xs.items():
+                xs[key] = val.numpy()
+            for key, val in preds.items():
+                preds[key] = val.detach.numpy()
+        signature = infer_signature(xs, preds)
+    else:
+        signature = None
+    writer.pytorch.log_model(model, artifact_path="model", signature=signature)
+
+
+TRACKER_MAPPING: dict[ExperimentTrackingTool, Callable[..., None]] = {
+    ExperimentTrackingTool.TENSORBOARD: write_tensorboard,
+    ExperimentTrackingTool.MLFLOW: write_mlflow,
+}
+
+LOGGER_MAPPING: dict[ExperimentTrackingTool, Callable[..., None]] = {
+    ExperimentTrackingTool.TENSORBOARD: log_hyperparams_tensorboard,
+    ExperimentTrackingTool.MLFLOW: log_hyperparams_mlflow,
+}
+
+PLOTTER_MAPPING: dict[ExperimentTrackingTool, Callable[..., None]] = {
+    ExperimentTrackingTool.TENSORBOARD: plot_tensorboard,
+    ExperimentTrackingTool.MLFLOW: plot_mlflow,
+}
+
+MODEL_LOGGER_MAPPING: dict[ExperimentTrackingTool, Callable[..., None]] = {
+    ExperimentTrackingTool.TENSORBOARD: log_model_tensorboard,
+    ExperimentTrackingTool.MLFLOW: log_model_mlflow,
+}
+
+
+def write_tracker(
+    *args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+) -> None:
+    return TRACKER_MAPPING[tracking_tool](*args)
+
+
+def log_tracker(
+    *args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+) -> None:
+    return LOGGER_MAPPING[tracking_tool](*args)
+
+
+def plot_tracker(
+    *args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+) -> None:
+    return PLOTTER_MAPPING[tracking_tool](*args)
+
+
+def log_model_tracker(
+    *args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+) -> None:
+    return MODEL_LOGGER_MAPPING[tracking_tool](*args)

{qadence-1.7.2 → qadence-1.7.3}/qadence/ml_tools/train_grad.py

@@ -1,10 +1,17 @@
 from __future__ import annotations
 
+import importlib
 import math
 from logging import getLogger
 from typing import Callable, Union
 
-from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn, TimeRemainingColumn
+from rich.progress import (
+    BarColumn,
+    Progress,
+    TaskProgressColumn,
+    TextColumn,
+    TimeRemainingColumn,
+)
 from torch import complex128, float32, float64
 from torch import device as torch_device
 from torch import dtype as torch_dtype
@@ -16,8 +23,15 @@ from torch.utils.tensorboard import SummaryWriter
 from qadence.ml_tools.config import TrainConfig
 from qadence.ml_tools.data import DictDataLoader, data_to_device
 from qadence.ml_tools.optimize_step import optimize_step
-from qadence.ml_tools.printing import print_metrics, write_tensorboard
+from qadence.ml_tools.printing import (
+    log_model_tracker,
+    log_tracker,
+    plot_tracker,
+    print_metrics,
+    write_tracker,
+)
 from qadence.ml_tools.saveload import load_checkpoint, write_checkpoint
+from qadence.types import ExperimentTrackingTool
 
 logger = getLogger(__name__)
 
@@ -30,7 +44,6 @@ def train(
     loss_fn: Callable,
     device: torch_device = None,
     optimize_step: Callable = optimize_step,
-    write_tensorboard: Callable = write_tensorboard,
     dtype: torch_dtype = None,
 ) -> tuple[Module, Optimizer]:
     """Runs the training loop with gradient-based optimizer.
@@ -48,15 +61,11 @@ def train(
             the model
         optimizer: The optimizer to use.
         config: `TrainConfig` with additional training options.
-        loss_fn: Loss function returning (loss: float, metrics: dict[str, float])
+        loss_fn: Loss function returning (loss: float, metrics: dict[str, float], ...)
         device: String defining device to train on, pass 'cuda' for GPU.
         optimize_step: Customizable optimization callback which is called at every iteration.=
             The function must have the signature `optimize_step(model,
             optimizer, loss_fn, xs, device="cpu")`.
-        write_tensorboard: Customizable tensorboard logging callback which is
-            called every `config.write_every` iterations. The function must have
-            the signature `write_tensorboard(writer, loss, metrics, iteration)`
-            (see the example below).
         dtype: The dtype to use for the data.
 
     Example:
@@ -122,8 +131,11 @@ def train(
         model = model.module.to(device=device, dtype=dtype)
     else:
         model = model.to(device=device, dtype=dtype)
-    # initialize tensorboard
-    writer = SummaryWriter(config.folder, purge_step=init_iter)
+    # initialize tracking tool
+    if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
+        writer = SummaryWriter(config.folder, purge_step=init_iter)
+    else:
+        writer = importlib.import_module("mlflow")
 
     perform_val = isinstance(config.val_every, int)
     if perform_val:
@@ -166,7 +178,7 @@ def train(
                 best_val_loss, metrics = loss_fn(model, xs_to_device)
 
                 metrics["val_loss"] = best_val_loss
-                write_tensorboard(writer, None, metrics, init_iter)
+                write_tracker(writer, None, metrics, init_iter, tracking_tool=config.tracking_tool)
 
             if config.folder:
                 if config.checkpoint_best_only:
@@ -174,6 +186,14 @@ def train(
                 else:
                     write_checkpoint(config.folder, model, optimizer, init_iter)
 
+            plot_tracker(
+                writer,
+                model,
+                init_iter,
+                config.plotting_functions,
+                tracking_tool=config.tracking_tool,
+            )
+
         except KeyboardInterrupt:
             logger.info("Terminating training gracefully after the current iteration.")
 
@@ -218,19 +238,31 @@ def train(
                     print_metrics(loss, metrics, iteration - 1)
 
                 if iteration % config.write_every == 0:
-                    write_tensorboard(writer, loss, metrics, iteration - 1)
+                    write_tracker(
+                        writer, loss, metrics, iteration, tracking_tool=config.tracking_tool
+                    )
 
+                if iteration % config.plot_every == 0:
+                    plot_tracker(
+                        writer,
+                        model,
+                        iteration,
+                        config.plotting_functions,
+                        tracking_tool=config.tracking_tool,
+                    )
                 if perform_val:
                     if iteration % config.val_every == 0:
                         xs = next(dl_iter_val)
                         xs_to_device = data_to_device(xs, device=device, dtype=data_dtype)
-                        val_loss, _ = loss_fn(model, xs_to_device)
+                        val_loss, *_ = loss_fn(model, xs_to_device)
                         if config.validation_criterion(val_loss, best_val_loss, config.val_epsilon):  # type: ignore[misc]
                             best_val_loss = val_loss
                             if config.folder and config.checkpoint_best_only:
                                 write_checkpoint(config.folder, model, optimizer, iteration="best")
                             metrics["val_loss"] = val_loss
-                            write_tensorboard(writer, None, metrics, iteration)
+                            write_tracker(
+                                writer, loss, metrics, iteration, tracking_tool=config.tracking_tool
+                            )
 
                 if config.folder:
                     if iteration % config.checkpoint_every == 0 and not config.checkpoint_best_only:
@@ -245,17 +277,32 @@ def train(
         try:
             xs = next(dl_iter) if dataloader is not None else None  # type: ignore[arg-type]
             xs_to_device = data_to_device(xs, device=device, dtype=data_dtype)
-            loss, metrics = loss_fn(model, xs_to_device)
+            loss, metrics, *_ = loss_fn(model, xs_to_device)
+            if dataloader is None:
+                loss = loss.item()
             if iteration % config.print_every == 0 and config.verbose:
                 print_metrics(loss, metrics, iteration)
 
         except KeyboardInterrupt:
             logger.info("Terminating training gracefully after the current iteration.")
 
-    # Final printing, writing and checkpointing
+    # Final checkpointing and writing
     if config.folder and not config.checkpoint_best_only:
         write_checkpoint(config.folder, model, optimizer, iteration)
-    write_tensorboard(writer, loss, metrics, iteration)
-    writer.close()
+    write_tracker(writer, loss, metrics, iteration, tracking_tool=config.tracking_tool)
+
+    # writing hyperparameters
+    if config.hyperparams:
+        log_tracker(writer, config.hyperparams, metrics, tracking_tool=config.tracking_tool)
+
+    # logging the model
+    if config.log_model:
+        log_model_tracker(writer, model, dataloader, tracking_tool=config.tracking_tool)
+
+    # close tracker
+    if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:
+        writer.close()
+    elif config.tracking_tool == ExperimentTrackingTool.MLFLOW:
+        writer.end_run()
 
     return model, optimizer