Repository: orobix/quadra (language: python)

Issue:
fix: Fix circular import making code not running properly
## Summary

Fix circular import introduced with latest pr merge even tough tests were passing properly.

## Type of Change

Please select the one relevant option below:

- Bug fix (non-breaking change that solves an issue)

## Checklist

Please confirm that the following tasks have been completed:

- [X] I have tested my changes locally and they work as expected. (Please describe the tests you performed.)
- [ ] I have added unit tests for my changes, or updated existing tests if necessary.
- [ ] I have updated the documentation, if applicable.
- [X] I have installed pre-commit and run locally for my code changes.



Discussion / hints:

Fix patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 03233dec..f7d371e4 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -5,8 +5,9 @@ repos:
     hooks:
       - id: ruff
         name: ruff
+        files: \.py$
         entry: ruff
-        language: system
+        language: python
         exclude: ^(tests|scripts|notebooks|quadra/utils/tests)/
   - repo: local
     hooks:
diff --git a/docs/images/model_manager/model_compare.png b/docs/images/model_manager/model_compare.png
new file mode 100644
index 00000000..6cf4eecf
Binary files /dev/null and b/docs/images/model_manager/model_compare.png differ
diff --git a/docs/images/model_manager/model_description.png b/docs/images/model_manager/model_description.png
new file mode 100644
index 00000000..16af05d3
Binary files /dev/null and b/docs/images/model_manager/model_description.png differ
diff --git a/docs/images/model_manager/registered_model.png b/docs/images/model_manager/registered_model.png
new file mode 100644
index 00000000..766b5b33
Binary files /dev/null and b/docs/images/model_manager/registered_model.png differ
diff --git a/docs/images/model_manager/stage_model.png b/docs/images/model_manager/stage_model.png
new file mode 100644
index 00000000..392658da
Binary files /dev/null and b/docs/images/model_manager/stage_model.png differ
diff --git a/docs/tutorials/configurations.md b/docs/tutorials/configurations.md
index 45e6dd94..babfb3f9 100644
--- a/docs/tutorials/configurations.md
+++ b/docs/tutorials/configurations.md
@@ -175,10 +175,13 @@ model:
 task:
   lr_multiplier: 0.0
   run_test: True
-  export_type: [torchscript]
   report: True
   output:
     example: True
+  export_config:
+    types: [torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
 
 core:
   tag: "run"
diff --git a/docs/tutorials/devices_setup.md b/docs/tutorials/devices_setup.md
index dbcb7956..54556e65 100644
--- a/docs/tutorials/devices_setup.md
+++ b/docs/tutorials/devices_setup.md
@@ -36,13 +36,15 @@ For Sklearn based tasks there's generally a `device` field in the configuration
 ```yaml
 _target_: quadra.tasks.SklearnClassification
 device: "cuda:0"
-export_type: [torchscript]
 output:
   folder: "classification_experiment"
   save_backbone: false
   report: true
   example: true
   test_full_data: true
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
 ```
 
 You can change the device to `cpu` or a different cuda device depending on your needs.
diff --git a/docs/tutorials/examples/classification.md b/docs/tutorials/examples/classification.md
index 17c5e972..4c21705a 100644
--- a/docs/tutorials/examples/classification.md
+++ b/docs/tutorials/examples/classification.md
@@ -111,10 +111,13 @@ task:
   lr_multiplier: 0.0
   gradcam: true
   run_test: True
-  export_type: [torchscript, pytorch]
   report: True
   output:
     example: True
+  export_config:
+    types: [pytorch, torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
 
 core:
   tag: "run"
@@ -181,10 +184,13 @@ datamodule:
 task:
   gradcam: True # Enable gradcam computation during evaluation
   run_test: True # Perform test evaluation at the end of training
-  export_type: [torchscript, pytorch]
   report: True 
   output:
     example: True # Generate an example of concordants and discordants predictions for each class
+  export_config:
+    types: [pytorch, torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
    
 model:
   num_classes: 3 # This is very important
diff --git a/docs/tutorials/examples/multilabel_classification.md b/docs/tutorials/examples/multilabel_classification.md
index 9b2eeedb..cd9fe0b6 100644
--- a/docs/tutorials/examples/multilabel_classification.md
+++ b/docs/tutorials/examples/multilabel_classification.md
@@ -136,10 +136,13 @@ model:
 
 task:
   run_test: True # Perform test evaluation at the end of training
-  export_type: [torchscript]
   report: False 
   output:
     example: False 
+  export_config:
+    types: [torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
   
 logger:
   mlflow:
diff --git a/docs/tutorials/examples/segmentation.md b/docs/tutorials/examples/segmentation.md
index 280e9041..d7264a1f 100644
--- a/docs/tutorials/examples/segmentation.md
+++ b/docs/tutorials/examples/segmentation.md
@@ -167,11 +167,14 @@ backbone:
     classes: 4 # The total number of classes (background + foreground)
 
 task:
-  export_type: [torchscript] # Export model to torchscript for inference
   run_test: true # run test after training is completed
   report: false # allows to generate reports
   evaluate: # custom evaluation toggles
     analysis: false # Perform in depth analysis
+  export_config:
+    types: [torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
     
 datamodule:
   data_path: /path/to/the/dataset # change the path to the dataset
diff --git a/docs/tutorials/examples/sklearn_classification.md b/docs/tutorials/examples/sklearn_classification.md
index 66f9b543..7c14bc26 100644
--- a/docs/tutorials/examples/sklearn_classification.md
+++ b/docs/tutorials/examples/sklearn_classification.md
@@ -101,7 +101,10 @@ backbone:
     freeze: true
 
 task:
-  export_type: [torchscript, pytorch]
+  export_config:
+    types: [pytorch, torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
 
 core:
   tag: "run"
@@ -143,16 +146,19 @@ datamodule:
 
 task:
   device: cuda:0
-  export_type: [torchscript, pytorch]
   output:
     folder: classification_experiment
     save_backbone: true
     report: true
     example: true
     test_full_data: true
+  export_config:
+    types: [pytorch, torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
 ```
 
-This will train a logistic regression classifier using a resnet18 backbone, resizing the images to 224x224 and using a 5-fold cross validation. The `class_to_idx` parameter is used to map the class names to indexes, the indexes will be used to train the classifier. The `output` parameter is used to specify the output folder and the type of output to save. The `export_type` parameter can be used to export the model in different formats, at the moment only `torchscript` is supported.
+This will train a logistic regression classifier using a resnet18 backbone, resizing the images to 224x224 and using a 5-fold cross validation. The `class_to_idx` parameter is used to map the class names to indexes, the indexes will be used to train the classifier. The `output` parameter is used to specify the output folder and the type of output to save. The `export_config.types` parameter can be used to export the model in different formats, at the moment `torchscript` and `pytorch` are supported.
 Since `save_backbone` is set to true, the backbone (in torchscript format) will be saved along with the classifier. `test_full_data` is used to specify if a final test should be performed on all the data (after training on the training and validation datasets).
 
 ### Run
diff --git a/docs/tutorials/examples/sklearn_patch_classification.md b/docs/tutorials/examples/sklearn_patch_classification.md
index fda80c7b..35db4453 100644
--- a/docs/tutorials/examples/sklearn_patch_classification.md
+++ b/docs/tutorials/examples/sklearn_patch_classification.md
@@ -222,11 +222,14 @@ task:
   device: cuda:2
   output:
     folder: classification_patch_experiment
-    export_type: [torchscript]
     save_backbone: false
     report: true
     example: true
     reconstruction_method: major_voting
+  export_config:
+    types: [torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
 ```
 
 This will train a resnet18 model on the given dataset, using 256 as batch size and skipping the background class during training.
diff --git a/docs/tutorials/examples/ssl.md b/docs/tutorials/examples/ssl.md
index a2905821..836796e5 100644
--- a/docs/tutorials/examples/ssl.md
+++ b/docs/tutorials/examples/ssl.md
@@ -160,7 +160,7 @@ The output folder should contain the following entries:
 checkpoints  config_resolved.yaml  config_tree.txt  data  deployment_model  main.log
 ```
 
-The `checkpoints` folder contains the saved `pytorch` lightning checkpoints. The `data` folder contains a joblib version of the datamodule containing all parameters and dataset spits. The `deployment_model` folder contains the model ready for production in the format specified in the task `export_type` parameter (default `torchscript`). 
+The `checkpoints` folder contains the saved `pytorch` lightning checkpoints. The `data` folder contains a joblib version of the datamodule containing all parameters and dataset spits. The `deployment_model` folder contains the model ready for production in the format specified in the task `export_config.types` parameter (default `torchscript`). 
 
 ### Run (Advanced) - Changing transformations
 
diff --git a/docs/tutorials/model_management.md b/docs/tutorials/model_management.md
new file mode 100644
index 00000000..be3313ca
--- /dev/null
+++ b/docs/tutorials/model_management.md
@@ -0,0 +1,150 @@
+# Model Management
+
+## Overview
+
+Model Management is the process of maintaining and orchestrating the lifecycle of machine learning models. It involves steps such as model creation, training, evaluation, deployment, and monitoring, among other steps. Model Versioning, on the other hand, allows data scientists to track and manage different versions of models during the model's lifecycle. It's crucial for reproducibility, collaboration, and consistent model performance. Quadra, out-of-the-box, supports managing models and their versions through [`MlflowModelManager`][quadra.utils.model_manager.MlflowModelManager] wrapper class with the help of [Mlflow](https://mlflow.org/) library. 
+
+ [`MlflowModelManager`][quadra.utils.model_manager.MlflowModelManager] is an extension of [`AbstractModelManager`][quadra.utils.model_manager.AbstractModelManager] which serves as a blueprint for model managers and specifies the required methods for managing models. It includes methods such as:
+
+- Register the model
+- Retrieve latest version
+- Transition the model to a new stage
+- Delete the model 
+
+By defining an abstract class, we establish a common interface that can be implemented by different model managers besides Mlflow.
+
+
+## Example Usage
+
+In this section, we will create example project for segmentation task and use [`MlflowModelManager`][quadra.utils.model_manager.MlflowModelManager] to manage the production model. `Quadra` provides a toy example where you can train a segmentation model for Oxford-IIIT Pet Dataset. 
+
+!!!note
+    You can find a detailed explanation for customizing the segmentation task under [Segmentation Example](/tutorials/examples/segmentation.md) section.
+
+First of all, we need to run `Mlflow` server with artifact store. You can find the instructions for running `Mlflow` server [here](https://mlflow.org/docs/latest/tracking.html#mlflow-tracking-servers). Let's open a new terminal and run the following command:
+
+```bash
+mlflow server \
+--backend-store-uri sqlite:///mlflow.db \
+--default-artifact-root file:///tmp/mlflow \
+--host 0.0.0.0
+```
+
+Then, we can start training from different terminal window while `Mlflow` server is running:
+
+```bash
+MLFLOW_TRACKING_URI="http://localhost:5000" \
+quadra experiment=generic/oxford_pet/segmentation/smp \
+trainer.max_epochs=5 \
+core.name=cats_vs_dogs \
+backbone.model.arch=unet,unetplusplus \
+backbone.model.encoder_name=resnet18,resnet50 \
+--multirun
+```
+
+This command will train a segmentation model for 5 epochs and save experiments run under `cats_vs_dogs` experiment tab. It will run the same experiment with different backbones or segmentation model architectures. After all trainings are completed, we open up the `Mlflow` UI and see experiments run under `cats_vs_dogs` directory. Under the artifacts section of each run, we can find `Mlflow` model artifact under the `deployment_model` directory. Other folders contains metadata or reports about the experiment run.
+
+### Registering the Model
+
+After all runs are completed, we can register the best model with the following command:
+
+```python
+import os
+from quadra.utils.model_manager import MlflowModelManager
+os.environ["MLFLOW_TRACKING_URI"] = "http://localhost:5000"
+
+manager = MlflowModelManager()
+manager.register_best_model(experiment_name="cats_vs_dogs",
+                            # metric we want to use for selecting the best model
+                            metric="val_loss_epoch", 
+                            # model path under the artifact store
+                            model_path="deployment_model", 
+                            # the name of the model to be registered
+                            model_name="cvsd_model", 
+                            # optional tags
+                            tags={"type":"segmentation"}, 
+                            # metric sorting order for selecting the best model
+                            mode="min")  
+```
+
+<p align="center">
+    <figure>
+        <img src="/images/model_manager/registered_model.png" title="Registered Model">
+        <figcaption>Registered Model</figcaption>
+    </figure>
+</p>
+
+
+### Staging the Model
+
+After registering the model, we can transition the model to a new stage. We can transition the model to `staging` stage with the following command:
+
+```python
+manager.transition_model(model_name="cvsd_model",
+                        version=1,
+                        stage="staging",
+                        description="Staging Model for demo")
+```
+
+<p align="center">
+    <figure>
+        <img src="/images/model_manager/stage_model.png" title="Staged Model">
+        <figcaption>Staged Model</figcaption>
+    </figure>
+</p>
+
+
+### Update the Model
+
+Let's say we want to update the model with better performance. Let's train the one of the configuration with more epochs and register the model again:
+
+```bash
+MLFLOW_TRACKING_URI="http://localhost:5000" \
+quadra experiment=generic/oxford_pet/segmentation/smp \
+trainer.max_epochs=20 \
+core.name=cats_vs_dogs \
+backbone.model.arch=unetplusplus \
+backbone.model.encoder_name=resnet50 
+```
+
+With this new model, we have improved `val_loss_epoch` metric compared to the previous model.
+
+<p align="center">
+    <figure>
+        <img src="/images/model_manager/model_compare.png" title="Comparison of Runs">
+        <figcaption>Comparison of Runs</figcaption>
+    </figure>
+</p>
+
+Let's register the new model with the following command:
+
+```python
+manager.register_model(model_location="runs:/<run-id>/deployment_model/model.pt",
+                       model_name="cvsd_model",
+                       tags={"type":"smp"},
+                       description="Better model"
+                       )
+```
+
+and finally transition the model to `production` stage:
+
+```python
+manager.transition_model(model_name="cvsd_model",
+                        version=2,
+                        stage="production",
+                        description="Production Model for demo")
+```
+
+!!!note
+    When we visit the model page from Mlflow UI, we can see the model version and stage information under the `Versions` tab. Moreover, we can see the model history automatically stored under the `Description` field.
+
+    <p align="center">
+        <figure>
+            <img src="/images/model_manager/model_description.png" title="Model Transition History">
+            <figcaption>Model Transition Historys</figcaption>
+        </figure>
+    </p>
+
+
+
+
diff --git a/docs/tutorials/tasks.md b/docs/tutorials/tasks.md
index 6be7d6e0..18a0eb88 100644
--- a/docs/tutorials/tasks.md
+++ b/docs/tutorials/tasks.md
@@ -50,4 +50,4 @@ graph LR;
 Most of these tasks have an associated evaluation task used for inference.
 ## Adding new tasks
 
-If your require the `pytorch-lightning` library, you can add a new task by extending the [`LightningTask`][quadra.base.LightningTask] class. Otherwise, you can simply start implementing a new task by extending the `Task` class.
\ No newline at end of file
+If you require the `pytorch-lightning` library, you can add a new task by extending the [`LightningTask`][quadra.tasks.base.LightningTask] class. Otherwise, you can simply start implementing a new task by extending the `Task` class.
\ No newline at end of file
diff --git a/mkdocs.yml b/mkdocs.yml
index fcc24a3d..ae89f3d2 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -14,6 +14,7 @@ nav:
       - Configurations: tutorials/configurations.md
       - Setting up devices: tutorials/devices_setup.md
       - Reproducibility: tutorials/reproducibility.md
+      - Model Management: tutorials/model_management.md
       - External projects integration: tutorials/integration.md
       - Contributing: tutorials/contribution.md
       - Building documentation: tutorials/documentation.md
diff --git a/pyproject.toml b/pyproject.toml
index 4b4bc7b5..cfa3b6ee 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -63,7 +63,8 @@ dependencies = [
   "h5py==3.8.*",
   "timm==0.6.12", # required by smp
   "segmentation-models-pytorch==0.3.*",
-  "anomalib@git+https://github.com/orobix/anomalib.git@v0.4.0+obx.1.0.1"
+  "anomalib@git+https://github.com/orobix/anomalib.git@v0.4.0+obx.1.0.1",
+  "xxhash==3.2.*",
 ]
 
 [project.optional-dependencies]
diff --git a/quadra/callbacks/mlflow.py b/quadra/callbacks/mlflow.py
index c77f9adc..17b7b6f4 100644
--- a/quadra/callbacks/mlflow.py
+++ b/quadra/callbacks/mlflow.py
@@ -9,6 +9,8 @@
 from pytorch_lightning.utilities import rank_zero_only
 from pytorch_lightning.utilities.types import STEP_OUTPUT
 
+from quadra.utils.mlflow import get_mlflow_logger
+
 
 def check_minio_credentials() -> None:
     """Check minio credentials for aws based storage such as minio.
@@ -41,38 +43,28 @@ def check_file_server_dependencies() -> None:
         ) from e
 
 
-def set_minio_credentials(access_key_id: str, secret_access_key: str) -> None:
-    """Setting minio credential as environment variables.
+def validate_artifact_storage(logger: MLFlowLogger):
+    """Validate artifact storage.
 
     Args:
-        access_key_id: string id
-        secret_access_key: string key.
+        logger: Mlflow logger from pytorch lightning.
 
-    Returns:
-        None
     """
-    os.environ["AWS_ACCESS_KEY_ID"] = access_key_id
-    os.environ["AWS_SECRET_ACCESS_KEY"] = secret_access_key
-
+    from quadra.utils.utils import get_logger  # pylint: disable=[import-outside-toplevel]
 
-def get_mlflow_logger(trainer: Trainer) -> Optional[MLFlowLogger]:
-    """Safely get Mlflow logger from Trainer loggers.
+    log = get_logger(__name__)
 
-    Args:
-        trainer: Pytorch Lightning trainer.
-
-    Returns:
-        An mlflow logger if available, else None.
-    """
-    if isinstance(trainer.logger, MLFlowLogger):
-        return trainer.logger
-
-    if isinstance(trainer.logger, list):
-        for logger in trainer.logger:
-            if isinstance(logger, MLFlowLogger):
-                return logger
-
-    return None
+    client = logger.experiment
+    # TODO: we have to access the internal api to get the artifact uri, however there could be a better way
+    artifact_uri = client._tracking_client._get_artifact_repo(  # pylint: disable=protected-access
+        logger.run_id
+    ).artifact_uri
+    if artifact_uri.startswith("s3://"):
+        check_minio_credentials()
+        check_file_server_dependencies()
+        log.info("Mlflow artifact storage is AWS/S3 basedand credentials and dependencies are satisfied.")
+    else:
+        log.info("Mlflow artifact storage uri is %s. Validation checks are not implemented.", artifact_uri)
 
 
 class UploadCodeAsArtifact(Callback):
@@ -88,7 +80,6 @@ class UploadCodeAsArtifact(Callback):
 
     def __init__(self, source_dir: str):
         self.source_dir = source_dir
-        check_file_server_dependencies()
 
     @rank_zero_only
     def on_test_end(self, trainer: Trainer, pl_module: LightningModule):
@@ -220,7 +211,6 @@ def __init__(
         self.ckpt_ext = ckpt_ext
         self.delete_after_upload = delete_after_upload
         self.upload = upload
-        check_file_server_dependencies()
 
     @rank_zero_only
     def on_test_end(self, trainer: Trainer, pl_module: LightningModule):
diff --git a/quadra/configs/callbacks/all.yaml b/quadra/configs/callbacks/all.yaml
index 4a6fae16..f8e75071 100644
--- a/quadra/configs/callbacks/all.yaml
+++ b/quadra/configs/callbacks/all.yaml
@@ -11,7 +11,7 @@ model_checkpoint:
 
 upload_code_as_artifact:
   _target_: quadra.callbacks.mlflow.UploadCodeAsArtifact
-  source_dir: ${work_dir}/src/
+  source_dir: ${work_dir}/quadra/
 
 upload_ckpts_as_artifact:
   _target_: quadra.callbacks.mlflow.UploadCheckpointsAsArtifact
diff --git a/quadra/configs/experiment/base/classification/classification.yaml b/quadra/configs/experiment/base/classification/classification.yaml
index da187780..3ce38c9b 100644
--- a/quadra/configs/experiment/base/classification/classification.yaml
+++ b/quadra/configs/experiment/base/classification/classification.yaml
@@ -24,7 +24,9 @@ model:
 task:
   lr_multiplier: 0.0
   run_test: True
-  export_type: [torchscript, pytorch]
+  export_config:
+    types: [torchscript, pytorch]
+    input_shapes:
   report: True
   gradcam: True
   output:
diff --git a/quadra/configs/experiment/base/classification/sklearn_classification.yaml b/quadra/configs/experiment/base/classification/sklearn_classification.yaml
index d8073cbe..1a833b99 100644
--- a/quadra/configs/experiment/base/classification/sklearn_classification.yaml
+++ b/quadra/configs/experiment/base/classification/sklearn_classification.yaml
@@ -14,8 +14,10 @@ backbone:
     freeze: true
 
 task:
-  export_type: [torchscript, pytorch]
-  
+  export_config:
+    types: [torchscript, pytorch]
+    input_shapes: null
+
 core:
   tag: "run"
   name: "sklearn-classification"
diff --git a/quadra/configs/logger/mlflow.yaml b/quadra/configs/logger/mlflow.yaml
index 812ff0b5..b59e22c9 100644
--- a/quadra/configs/logger/mlflow.yaml
+++ b/quadra/configs/logger/mlflow.yaml
@@ -9,3 +9,4 @@ mlflow:
   prefix: ""
   artifact_location: null
   run_id: null
+  log_model: false
diff --git a/quadra/configs/task/anomalib/cfa.yaml b/quadra/configs/task/anomalib/cfa.yaml
index 8099f704..f675715c 100644
--- a/quadra/configs/task/anomalib/cfa.yaml
+++ b/quadra/configs/task/anomalib/cfa.yaml
@@ -1,6 +1,9 @@
 _target_: quadra.tasks.anomaly.AnomalibDetection
 report: true
 run_test: true
-export_type: [torchscript]
 module_function:
   _target_: anomalib.models.cfa.lightning_model.CfaLightning
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
+
diff --git a/quadra/configs/task/anomalib/cflow.yaml b/quadra/configs/task/anomalib/cflow.yaml
index ebf7d496..9cd0912b 100644
--- a/quadra/configs/task/anomalib/cflow.yaml
+++ b/quadra/configs/task/anomalib/cflow.yaml
@@ -1,6 +1,9 @@
 _target_: quadra.tasks.anomaly.AnomalibDetection
 report: true
 run_test: true
-export_type: null
 module_function:
   _target_: anomalib.models.cflow.lightning_model.CflowLightning
+export_config:
+  types:
+    - torchscript
+  input_shapes: null
diff --git a/quadra/configs/task/anomalib/csflow.yaml b/quadra/configs/task/anomalib/csflow.yaml
index e6219977..68a1ec24 100644
--- a/quadra/configs/task/anomalib/csflow.yaml
+++ b/quadra/configs/task/anomalib/csflow.yaml
@@ -1,6 +1,9 @@
 _target_: quadra.tasks.anomaly.AnomalibDetection
 report: true
 run_test: true
-export_type: [torchscript]
 module_function:
   _target_: anomalib.models.csflow.lightning_model.CsflowLightning
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
+
diff --git a/quadra/configs/task/anomalib/draem.yaml b/quadra/configs/task/anomalib/draem.yaml
index 7783b6fe..09d29e51 100644
--- a/quadra/configs/task/anomalib/draem.yaml
+++ b/quadra/configs/task/anomalib/draem.yaml
@@ -1,6 +1,9 @@
 _target_: quadra.tasks.anomaly.AnomalibDetection
 report: true
 run_test: true
-export_type: [torchscript]
 module_function:
   _target_: anomalib.models.draem.lightning_model.DraemLightning
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
+
diff --git a/quadra/configs/task/anomalib/fastflow.yaml b/quadra/configs/task/anomalib/fastflow.yaml
index ada2599c..d681d71d 100644
--- a/quadra/configs/task/anomalib/fastflow.yaml
+++ b/quadra/configs/task/anomalib/fastflow.yaml
@@ -1,6 +1,9 @@
 _target_: quadra.tasks.anomaly.AnomalibDetection
 report: true
 run_test: true
-export_type: [torchscript]
 module_function:
   _target_: anomalib.models.fastflow.lightning_model.FastflowLightning
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
+
diff --git a/quadra/configs/task/anomalib/padim.yaml b/quadra/configs/task/anomalib/padim.yaml
index 929eca85..25deb46d 100644
--- a/quadra/configs/task/anomalib/padim.yaml
+++ b/quadra/configs/task/anomalib/padim.yaml
@@ -1,6 +1,9 @@
 _target_: quadra.tasks.anomaly.AnomalibDetection
 report: true
 run_test: true
-export_type: [torchscript]
 module_function:
   _target_: anomalib.models.padim.lightning_model.PadimLightning
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
+
diff --git a/quadra/configs/task/anomalib/patchcore.yaml b/quadra/configs/task/anomalib/patchcore.yaml
index e6945c44..4dfdc376 100644
--- a/quadra/configs/task/anomalib/patchcore.yaml
+++ b/quadra/configs/task/anomalib/patchcore.yaml
@@ -1,6 +1,9 @@
 _target_: quadra.tasks.anomaly.AnomalibDetection
 report: true
 run_test: true
-export_type: [torchscript]
 module_function:
   _target_: anomalib.models.patchcore.lightning_model.PatchcoreLightning
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
+
diff --git a/quadra/configs/task/classification.yaml b/quadra/configs/task/classification.yaml
index 0974aac8..a2ecb8e3 100644
--- a/quadra/configs/task/classification.yaml
+++ b/quadra/configs/task/classification.yaml
@@ -1,7 +1,9 @@
 _target_: quadra.tasks.Classification
-export_type: [torchscript, pytorch]
 lr_multiplier: null
 output:
   example: false
 report: false
 run_test: true
+export_config:
+  types: [pytorch, torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
diff --git a/quadra/configs/task/segmentation.yaml b/quadra/configs/task/segmentation.yaml
index a15bc98e..51b0b2e3 100644
--- a/quadra/configs/task/segmentation.yaml
+++ b/quadra/configs/task/segmentation.yaml
@@ -1,10 +1,12 @@
 _target_: quadra.tasks.Segmentation
 num_viz_samples: 10
 checkpoint_path: null
-export_type: [torchscript]
 run_test: true
 report: false
 evaluate:
   analysis: false
   control: false
   corruption: false
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
diff --git a/quadra/configs/task/sklearn_classification.yaml b/quadra/configs/task/sklearn_classification.yaml
index c904377f..feae7377 100644
--- a/quadra/configs/task/sklearn_classification.yaml
+++ b/quadra/configs/task/sklearn_classification.yaml
@@ -1,9 +1,11 @@
 _target_: quadra.tasks.SklearnClassification
 device: "cuda:0"
-export_type: [torchscript, pytorch]
 output:
   folder: "classification_experiment"
   save_backbone: false
   report: true
   example: true
   test_full_data: true
+export_config:
+  types: [pytorch, torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
diff --git a/quadra/configs/task/sklearn_classification_patch.yaml b/quadra/configs/task/sklearn_classification_patch.yaml
index b38a1d7a..2a25c366 100644
--- a/quadra/configs/task/sklearn_classification_patch.yaml
+++ b/quadra/configs/task/sklearn_classification_patch.yaml
@@ -1,9 +1,11 @@
 _target_: quadra.tasks.PatchSklearnClassification
 device: cuda:2
-export_type: [torchscript]
 output:
   folder: classification_patch_experiment
   save_backbone: false
   report: true
   example: true
   reconstruction_method: major_voting
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
diff --git a/quadra/configs/task/ssl.yaml b/quadra/configs/task/ssl.yaml
index cc32f5db..24ae88e4 100644
--- a/quadra/configs/task/ssl.yaml
+++ b/quadra/configs/task/ssl.yaml
@@ -1,3 +1,5 @@
 _target_: quadra.tasks.SSL
-export_type: [torchscript]
 run_test: false
+export_config:
+  types: [torchscript]
+  input_shapes:
diff --git a/quadra/datamodules/base.py b/quadra/datamodules/base.py
index 7b6796be..ac284d2e 100644
--- a/quadra/datamodules/base.py
+++ b/quadra/datamodules/base.py
@@ -1,18 +1,22 @@
+import multiprocessing as mp
+import multiprocessing.pool as mpp
 import os
 import pickle as pkl
+import typing
 from functools import wraps
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
+from typing import Any, Callable, Dict, Iterable, List, Literal, Optional, Sequence, Tuple, Union, cast
 
 import albumentations
 import numpy as np
 import pandas as pd
 import torch
+import xxhash
 from pytorch_lightning import LightningDataModule
+from tqdm import tqdm
 
 from quadra.utils import utils
 
 log = utils.get_logger(__name__)
-IMAGE_EXTENSIONS: List[str] = [".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".pbm", ".pgm", ".ppm", ".pxm", ".pnm"]
 TrainDataset = Union[torch.utils.data.Dataset, Sequence[torch.utils.data.Dataset]]
 ValDataset = Union[torch.utils.data.Dataset, Sequence[torch.utils.data.Dataset]]
 TestDataset = torch.utils.data.Dataset
@@ -46,6 +50,73 @@ def __new__(cls, name, bases, dct):
         return super().__new__(cls, name, bases, dct)
 
 
+def compute_file_content_hash(path: str, hash_size: Literal[32, 64, 128] = 64) -> str:
+    """Get hash of a file based on its content.
+
+    Args:
+        path: Path to the file.
+        hash_size: Size of the hash. Must be one of [32, 64, 128].
+
+    Returns:
+        The hash of the file.
+    """
+    with open(path, "rb") as f:
+        data = f.read()
+
+        if hash_size == 32:
+            file_hash = xxhash.xxh32(data, seed=42).hexdigest()
+        elif hash_size == 64:
+            file_hash = xxhash.xxh64(data, seed=42).hexdigest()
+        elif hash_size == 128:
+            file_hash = xxhash.xxh128(data, seed=42).hexdigest()
+        else:
+            raise ValueError(f"Invalid hash size {hash_size}. Must be one of [32, 64, 128].")
+
+    return file_hash
+
+
+def compute_file_size_hash(path: str, hash_size: Literal[32, 64, 128] = 64) -> str:
+    """Get hash of a file based on its size.
+
+    Args:
+        path: Path to the file.
+        hash_size: Size of the hash. Must be one of [32, 64, 128].
+
+    Returns:
+        The hash of the file.
+    """
+    data = str(os.path.getsize(path))
+
+    if hash_size == 32:
+        file_hash = xxhash.xxh32(data, seed=42).hexdigest()
+    elif hash_size == 64:
+        file_hash = xxhash.xxh64(data, seed=42).hexdigest()
+    elif hash_size == 128:
+        file_hash = xxhash.xxh128(data, seed=42).hexdigest()
+    else:
+        raise ValueError(f"Invalid hash size {hash_size}. Must be one of [32, 64, 128].")
+
+    return file_hash
+
+
+@typing.no_type_check
+def istarmap(self, func: Callable, iterable: Iterable, chunksize: int = 1):
+    # pylint: disable=all
+    """Starmap-version of imap."""
+    self._check_running()
+    if chunksize < 1:
+        raise ValueError("Chunksize must be 1+, not {0:n}".format(chunksize))
+
+    task_batches = mpp.Pool._get_tasks(func, iterable, chunksize)
+    result = mpp.IMapIterator(self)
+    self._taskqueue.put((self._guarded_task_generation(result._job, mpp.starmapstar, task_batches), result._set_length))
+    return (item for chunk in result for item in chunk)
+
+
+# Patch Pool class to include istarmap
+mpp.Pool.istarmap = istarmap  # type: ignore[attr-defined]
+
+
 class BaseDataModule(LightningDataModule, metaclass=DecorateParentMethod):
     """Base class for all data modules.
 
@@ -61,6 +132,10 @@ class BaseDataModule(LightningDataModule, metaclass=DecorateParentMethod):
             Defaults to None.
         test_transform: Transformations for test dataset.
             Defaults to None.
+        enable_hashing: Whether to enable hashing of images. Defaults to True.
+        hash_size: Size of the hash. Must be one of [32, 64, 128]. Defaults to 64.
+        hash_type: Type of hash to use, if content hash is used, the hash is computed on the file content, otherwise
+            the hash is computed on the file size which is faster but less safe. Defaults to "content".
     """
 
     def __init__(
@@ -78,6 +153,9 @@ def __init__(
         train_transform: Optional[albumentations.Compose] = None,
         val_transform: Optional[albumentations.Compose] = None,
         test_transform: Optional[albumentations.Compose] = None,
+        enable_hashing: bool = True,
+        hash_size: Literal[32, 64, 128] = 64,
+        hash_type: Literal["content", "size"] = "content",
     ):
         super().__init__()
         self.num_workers = num_workers
@@ -88,6 +166,12 @@ def __init__(
         self.train_transform = train_transform
         self.val_transform = val_transform
         self.test_transform = test_transform
+        self.enable_hashing = enable_hashing
+        self.hash_size = hash_size
+        self.hash_type = hash_type
+
+        if self.hash_size not in [32, 64, 128]:
+            raise ValueError(f"Invalid hash size {self.hash_size}. Must be one of [32, 64, 128].")
 
         self.load_aug_images = load_aug_images
         self.aug_name = aug_name
@@ -102,6 +186,7 @@ def __init__(
         self.data_folder = "data"
         os.makedirs(self.data_folder, exist_ok=True)
         self.datamodule_checkpoint_file = os.path.join(self.data_folder, "datamodule.pkl")
+        self.dataset_file = os.path.join(self.data_folder, "dataset.csv")
 
     @property
     def train_data(self) -> pd.DataFrame:
@@ -166,12 +251,35 @@ def _prepare_data(self) -> None:
             "contained in the prepare_data method of a LightningModule."
         )
 
+    def hash_data(self) -> None:
+        """Computes the hash of the files inside the datasets."""
+        if not self.enable_hashing:
+            return
+
+        # TODO: We need to find a way to annotate the columns of data.
+        paths_and_hash_length = zip(self.data["samples"], [self.hash_size] * len(self.data))
+
+        with mp.Pool(min(8, mp.cpu_count() - 1)) as pool:
+            self.data["hash"] = list(
+                tqdm(
+                    pool.istarmap(  # type: ignore[attr-defined]
+                        compute_file_content_hash if self.hash_type == "content" else compute_file_size_hash,
+                        paths_and_hash_length,
+                    ),
+                    total=len(self.data),
+                    desc="Computing hashes",
+                )
+            )
+
+        self.data["hash_type"] = self.hash_type
+
     def prepare_data(self) -> None:
         """Prepares the data, should be overridden by subclasses."""
         if hasattr(self, "data"):
             return
 
         self._prepare_data()
+        self.hash_data()
         self.save_checkpoint()
 
     def __getstate__(self) -> Dict[str, Any]:
@@ -190,13 +298,11 @@ def save_checkpoint(self) -> None:
         datamodule to disk because we can't assign attributes to the datamodule in prepare_data when working with
         multiple gpus.
         """
-        saved_to_disk = False
-        if not os.path.exists(self.datamodule_checkpoint_file):
+        if not os.path.exists(self.datamodule_checkpoint_file) and not os.path.exists(self.dataset_file):
             with open(self.datamodule_checkpoint_file, "wb") as f:
                 pkl.dump(self, f)
-            saved_to_disk = True
 
-        if saved_to_disk:
+            self.data.to_csv(self.dataset_file, index=False)
             log.info("Datamodule checkpoint saved to disk.")
 
         if "targets" in self.data and not isinstance(self.data["targets"].iloc[0], np.ndarray):
diff --git a/quadra/datamodules/classification.py b/quadra/datamodules/classification.py
index 399cc903..6dfcebc4 100644
--- a/quadra/datamodules/classification.py
+++ b/quadra/datamodules/classification.py
@@ -12,7 +12,7 @@
 from timm.data.parsers.parser_image_folder import find_images_and_targets
 from torch.utils.data import DataLoader
 
-from quadra.datamodules.base import IMAGE_EXTENSIONS, BaseDataModule
+from quadra.datamodules.base import BaseDataModule
 from quadra.datasets import ImageClassificationListDataset
 from quadra.datasets.classification import MultilabelClassificationDataset
 from quadra.utils import utils
@@ -135,7 +135,7 @@ def _find_images_and_targets(
     ) -> Tuple[List[Tuple[str, int]], Dict[str, int]]:
         """Collects the samples from item folders."""
         images_and_targets, class_to_idx = find_images_and_targets(
-            folder=root_folder, types=IMAGE_EXTENSIONS, class_to_idx=class_to_idx
+            folder=root_folder, types=utils.IMAGE_EXTENSIONS, class_to_idx=class_to_idx
         )
         return images_and_targets, class_to_idx
 
@@ -397,6 +397,7 @@ class SklearnClassificationDataModule(BaseDataModule):
         label_map: Dictionary of conversion btw folder name and label.
         train_split_file: Optional path to a csv file containing the train split samples.
         test_split_file: Optional path to a csv file containing the test split samples.
+        **kwargs: Additional arguments for BaseDataModule
     """
 
     def __init__(
@@ -422,6 +423,7 @@ def __init__(
         test_split_file: Optional[str] = None,
         name: str = "sklearn_classification_datamodule",
         dataset: Type[ImageClassificationListDataset] = ImageClassificationListDataset,
+        **kwargs: Any,
     ):
         super().__init__(
             data_path=data_path,
@@ -432,6 +434,7 @@ def __init__(
             train_transform=train_transform,
             val_transform=val_transform,
             test_transform=test_transform,
+            **kwargs,
         )
 
         self.class_to_idx = class_to_idx
diff --git a/quadra/datamodules/generic/oxford_pet.py b/quadra/datamodules/generic/oxford_pet.py
index fbd66717..b52342a7 100644
--- a/quadra/datamodules/generic/oxford_pet.py
+++ b/quadra/datamodules/generic/oxford_pet.py
@@ -1,5 +1,5 @@
 import os
-from typing import Optional, Type
+from typing import Any, Optional, Type
 
 import albumentations
 import cv2
@@ -44,6 +44,7 @@ def __init__(
         train_transform: Optional[albumentations.Compose] = None,
         test_transform: Optional[albumentations.Compose] = None,
         val_transform: Optional[albumentations.Compose] = None,
+        **kwargs: Any,
     ):
         super().__init__(
             data_path=data_path,
@@ -57,6 +58,7 @@ def __init__(
             val_transform=val_transform,
             batch_size=batch_size,
             num_workers=num_workers,
+            **kwargs,
         )
 
     _RESOURCES = (
diff --git a/quadra/datamodules/segmentation.py b/quadra/datamodules/segmentation.py
index a9ae46fe..5fc5058a 100644
--- a/quadra/datamodules/segmentation.py
+++ b/quadra/datamodules/segmentation.py
@@ -2,7 +2,7 @@
 import glob
 import os
 import random
-from typing import Dict, List, Optional, Tuple, Type
+from typing import Any, Dict, List, Optional, Tuple, Type
 
 import albumentations
 import cv2
@@ -59,6 +59,7 @@ def __init__(
         val_split_file: Optional[str] = None,
         num_data_class: Optional[int] = None,
         exclude_good: bool = False,
+        **kwargs: Any,
     ):
         super().__init__(
             data_path=data_path,
@@ -69,6 +70,7 @@ def __init__(
             train_transform=train_transform,
             val_transform=val_transform,
             test_transform=test_transform,
+            **kwargs,
         )
         self.test_size = test_size
         self.val_size = val_size
@@ -413,6 +415,7 @@ def __init__(
         exclude_good: bool = False,
         num_data_train: Optional[int] = None,
         one_hot_encoding: bool = False,
+        **kwargs: Any,
     ):
         super().__init__(
             data_path=data_path,
@@ -423,6 +426,7 @@ def __init__(
             train_transform=train_transform,
             val_transform=val_transform,
             test_transform=test_transform,
+            **kwargs,
         )
         self.test_size = test_size
         self.val_size = val_size
diff --git a/quadra/datasets/anomaly.py b/quadra/datasets/anomaly.py
index ce02f10c..b1058fcd 100644
--- a/quadra/datasets/anomaly.py
+++ b/quadra/datasets/anomaly.py
@@ -11,7 +11,7 @@
 from torch import Tensor
 from torch.utils.data import Dataset
 
-from quadra.datamodules.base import IMAGE_EXTENSIONS
+from quadra.utils.utils import IMAGE_EXTENSIONS
 
 
 def create_validation_set_from_test_set(samples: DataFrame, seed: int = 0) -> DataFrame:
diff --git a/quadra/datasets/patch.py b/quadra/datasets/patch.py
index 001c0497..f4d34c9e 100644
--- a/quadra/datasets/patch.py
+++ b/quadra/datasets/patch.py
@@ -91,7 +91,7 @@ def __init__(
     def __getitem__(self, idx) -> Tuple[np.ndarray, np.ndarray]:
         path, y = self.samples[idx]
 
-        h5_file = h5py.File(os.path.join(self.data_path, "train", path))
+        h5_file = h5py.File(path)
         x = cv2.imread(os.path.join(self.data_path, h5_file["img_path"][()].decode("utf-8")))
 
         weights = h5_file["triangles_weights"][()]
diff --git a/quadra/models/base.py b/quadra/models/base.py
new file mode 100644
index 00000000..a733e881
--- /dev/null
+++ b/quadra/models/base.py
@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+import inspect
+from typing import Any, Sequence
+
+import torch
+from torch import nn
+
+
+class ModelSignatureWrapper(nn.Module):
+    """Model wrapper used to retrieve input shape. It can be used as a decorator of nn.Module, the first call to the
+    forward method will retrieve the input shape and store it in the input_shapes attribute.
+    """
+
+    def __init__(self, model: nn.Module):
+        super().__init__()
+        self.instance = model
+        self.input_shapes: Any = None
+        self.disable = False
+
+        if isinstance(self.instance, ModelSignatureWrapper):
+            # Handle nested ModelSignatureWrapper
+            self.input_shapes = self.instance.input_shapes
+            self.instance = self.instance.instance
+
+    def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        """Retrieve the input shape and forward the model."""
+        if self.input_shapes is None and not self.disable:
+            try:
+                self.input_shapes = self._get_input_shapes(*args, **kwargs)
+            except Exception:
+                # Avoid circular import
+                # pylint: disable=import-outside-toplevel
+                from quadra.utils.utils import get_logger  # noqa
+
+                log = get_logger(__name__)
+                log.warning(
+                    "Failed to retrieve input shapes after forward! To export the model you'll need to "
+                    "provide the input shapes manually setting the export_config.input_shapes parameter! "
+                    "Alternatively you could try to use a forward with supported input types (and their compositions) "
+                    "(list, tuple, dict, tensors)."
+                )
+                self.disable = True
+
+        return self.instance.forward(*args, **kwargs)
+
+    def to(self, *args, **kwargs):
+        """Handle calls to to method returning the underlying model."""
+        return ModelSignatureWrapper(self.instance.to(*args, **kwargs))
+
+    def _get_input_shapes(self, *args: Any, **kwargs: Any) -> list[Any]:
+        """Retrieve the input shapes from the input. Inputs will be in the same order as the forward method
+        signature.
+        """
+        input_shapes = []
+
+        for arg in args:
+            input_shapes.append(self._get_input_shape(arg))
+
+        signature = inspect.signature(self.instance.forward)
+
+        for i, key in enumerate(signature.parameters.keys()):
+            if i < len(args):
+                continue
+
+            if key in kwargs:
+                input_shapes.append(self._get_input_shape(kwargs[key]))
+            else:
+                # Retrieve the default value
+                input_shapes.append(self._get_input_shape(signature.parameters[key].default))
+
+        return input_shapes
+
+    def _get_input_shape(self, inp: Sequence | torch.Tensor) -> list[Any] | tuple[Any, ...] | dict[str, Any]:
+        """Recursive function to retrieve the input shapes."""
+        if isinstance(inp, list):
+            return [self._get_input_shape(i) for i in inp]
+
+        if isinstance(inp, tuple):
+            return tuple(self._get_input_shape(i) for i in inp)
+
+        if isinstance(inp, torch.Tensor):
+            return tuple(inp.shape[1:])
+
+        if isinstance(inp, dict):
+            return {k: self._get_input_shape(v) for k, v in inp.items()}
+
+        raise ValueError(f"Input type {type(inp)} not supported")
+
+    def __getattr__(self, name: str) -> torch.Tensor | nn.Module:
+        if name in ["instance", "input_shapes"]:
+            return self.__dict__[name]
+
+        return getattr(self.__dict__["instance"], name)
+
+    def __setattr__(self, name: str, value: torch.Tensor | nn.Module) -> None:
+        if name in ["instance", "input_shapes"]:
+            self.__dict__[name] = value
+        else:
+            setattr(self.instance, name, value)
+
+    def __getattribute__(self, __name: str) -> Any:
+        if __name in ["instance", "input_shapes", "__dict__", "forward", "_get_input_shapes", "_get_input_shape", "to"]:
+            return super().__getattribute__(__name)
+
+        return getattr(self.instance, __name)
diff --git a/quadra/modules/base.py b/quadra/modules/base.py
index 1f7ae004..ae030d26 100644
--- a/quadra/modules/base.py
+++ b/quadra/modules/base.py
@@ -8,6 +8,8 @@
 from torch import nn
 from torch.optim import Optimizer
 
+from quadra.models.base import ModelSignatureWrapper
+
 __all__ = ["BaseLightningModule", "SSLModule"]
 
 
@@ -28,7 +30,7 @@ def __init__(
         lr_scheduler_interval: Optional[str] = "epoch",
     ):
         super().__init__()
-        self.model = model
+        self.model = ModelSignatureWrapper(model)
         self.optimizer = optimizer
         self.schedulers = lr_scheduler
         self.lr_scheduler_interval = lr_scheduler_interval
@@ -90,7 +92,6 @@ def __init__(
         lr_scheduler: Optional[object] = None,
         lr_scheduler_interval: Optional[str] = "epoch",
     ):
-
         super().__init__(model, optimizer, lr_scheduler, lr_scheduler_interval)
         self.criterion = criterion
         self.classifier_train_loader: Optional[torch.utils.data.DataLoader]
@@ -176,7 +177,6 @@ def __init__(
         optimizer: Optional[Optimizer] = None,
         lr_scheduler: Optional[object] = None,
     ):
-
         super().__init__(model, optimizer, lr_scheduler)
         self.loss_fun = loss_fun
 
diff --git a/quadra/tasks/anomaly.py b/quadra/tasks/anomaly.py
index 5feee32c..3f36958a 100644
--- a/quadra/tasks/anomaly.py
+++ b/quadra/tasks/anomaly.py
@@ -20,6 +20,7 @@
 
 from quadra.callbacks.mlflow import get_mlflow_logger
 from quadra.datamodules import AnomalyDataModule
+from quadra.modules.base import ModelSignatureWrapper
 from quadra.tasks.base import Evaluation, LightningTask
 from quadra.utils import utils
 from quadra.utils.classification import get_results
@@ -40,7 +41,11 @@ class AnomalibDetection(Generic[AnomalyDataModuleT], LightningTask[AnomalyDataMo
             Defaults to None.
         run_test: Whether to run the test after training. Defaults to False.
         report: Whether to report the results. Defaults to False.
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
@@ -50,14 +55,14 @@ def __init__(
         checkpoint_path: Optional[str] = None,
         run_test: bool = True,
         report: bool = True,
-        export_type: Optional[List[str]] = None,
+        export_config: Optional[DictConfig] = None,
     ):
         super().__init__(
             config=config,
             checkpoint_path=checkpoint_path,
             run_test=run_test,
             report=report,
-            export_type=export_type,
+            export_config=export_config,
         )
         self._module: AnomalyModule
         self.module_function = module_function
@@ -103,10 +108,11 @@ def prepare(self) -> None:
         """Prepare the task."""
         super().prepare()
         self.module = self.config.model
+        self.module.model = ModelSignatureWrapper(self.module.model)
 
     def export(self) -> None:
         """Export model for production."""
-        if self.export_type is None or len(self.export_type) == 0:
+        if self.export_config is None or len(self.export_config.types) == 0:
             log.info("No export type specified skipping export")
             return
 
@@ -114,19 +120,29 @@ def export(self) -> None:
             log.warning("Skipping export since fast_dev_run is enabled")
             return
 
-        input_width = self.config.transforms.get("input_width")
-        input_height = self.config.transforms.get("input_height")
-
         model = self.module.model
 
-        half_precision = self.trainer.precision == 16
+        input_shapes = self.export_config.input_shapes
 
-        for export_type in self.export_type:
+        half_precision = int(self.trainer.precision) == 16
+
+        for export_type in self.export_config.types:
             if export_type == "torchscript":
-                export_torchscript_model(model, (1, 3, input_height, input_width), self.export_folder, half_precision)
+                out = export_torchscript_model(
+                    model=model,
+                    input_shapes=input_shapes,
+                    output_path=self.export_folder,
+                    half_precision=half_precision,
+                )
+
+                if out is None:
+                    log.warning("Skipping torchscript export since the model is not supported")
+                    continue
+
+                _, input_shapes = out
 
         model_json = {
-            "input_size": [self.config.transforms.input_width, self.config.transforms.input_height, 3],
+            "input_size": input_shapes,
             "classes": {0: "good", 1: "defect"},
             "mean": list(self.config.transforms.mean),
             "std": list(self.config.transforms.std),
diff --git a/quadra/tasks/base.py b/quadra/tasks/base.py
index f237239b..3b0f1ed5 100644
--- a/quadra/tasks/base.py
+++ b/quadra/tasks/base.py
@@ -8,13 +8,15 @@
 from hydra.core.hydra_config import HydraConfig
 from omegaconf import DictConfig, OmegaConf
 from pytorch_lightning import Callback, LightningModule, Trainer
-from pytorch_lightning.loggers import Logger
+from pytorch_lightning.loggers import Logger, MLFlowLogger
 from pytorch_lightning.utilities.device_parser import parse_gpu_ids
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from torch import nn
 from torch.jit._script import RecursiveScriptModule
+from torch.nn import Module
 
 from quadra import get_version
+from quadra.callbacks.mlflow import validate_artifact_storage
 from quadra.datamodules.base import BaseDataModule
 from quadra.utils import utils
 from quadra.utils.export import import_deployment_model
@@ -28,12 +30,17 @@ class Task(Generic[DataModuleT]):
 
     Args:
         config: The experiment configuration.
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
-    def __init__(self, config: DictConfig, export_type: Optional[List[str]] = None):
+    def __init__(self, config: DictConfig, export_config: Optional[DictConfig] = None):
         self.config = config
-        self.export_type = export_type
+        self.export_config = export_config
+        self.export_folder: str = "deployment_model"
         self._datamodule: DataModuleT
         self.metadata: Dict[str, Any]
         self.save_config()
@@ -85,7 +92,7 @@ def execute(self) -> None:
         self.prepare()
         self.train()
         self.test()
-        if self.export_type is not None and len(self.export_type) > 0:
+        if self.export_config is not None and len(self.export_config.types) > 0:
             self.export()
         self.generate_report()
         self.finalize()
@@ -99,7 +106,11 @@ class LightningTask(Generic[DataModuleT], Task[DataModuleT]):
         checkpoint_path: The path to the checkpoint to load the model from. Defaults to None.
         run_test: Whether to run the test after training. Defaults to False.
         report: Whether to generate a report. Defaults to False.
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
@@ -108,9 +119,9 @@ def __init__(
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
         report: bool = False,
-        export_type: Optional[List[str]] = None,
+        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config, export_type=export_type)
+        super().__init__(config, export_config=export_config)
         self.config = config
         self.checkpoint_path = checkpoint_path
         self.run_test = run_test
@@ -125,11 +136,13 @@ def prepare(self) -> None:
         """Prepare the experiment."""
         super().prepare()
 
+        # First setup loggers since some callbacks might need logger setup correctly.
+        if "logger" in self.config:
+            self.logger = self.config.logger
+
         if "callbacks" in self.config:
             self.callbacks = self.config.callbacks
 
-        if "logger" in self.config:
-            self.logger = self.config.logger
         self.devices = self.config.trainer.devices
         self.trainer = self.config.trainer
 
@@ -202,6 +215,8 @@ def logger(self, logger_config) -> None:
             if "_target_" in lg_conf:
                 log.info("Instantiating logger <%s>", lg_conf["_target_"])
                 logger = hydra.utils.instantiate(lg_conf)
+                if isinstance(logger, MLFlowLogger):
+                    validate_artifact_storage(logger)
                 instantiated_loggers.append(logger)
 
         self._logger = instantiated_loggers
@@ -236,6 +251,7 @@ def train(self) -> None:
             model=self.module,
             trainer=self.trainer,
         )
+
         self.trainer.fit(model=self.module, datamodule=self.datamodule)
 
     def test(self) -> Any:
@@ -248,11 +264,12 @@ def finalize(self) -> None:
         super().finalize()
         utils.finish(
             config=self.config,
-            model=self.module,
+            module=self.module,
             datamodule=self.datamodule,
             trainer=self.trainer,
             callbacks=self.callbacks,
             logger=self.logger,
+            export_folder=self.export_folder,
         )
 
         if not self.config.trainer.get("fast_dev_run"):
@@ -276,7 +293,7 @@ def execute(self) -> None:
         self.train()
         if self.run_test:
             self.test()
-        if self.export_type is not None and len(self.export_type) > 0:
+        if self.export_config is not None and len(self.export_config.types) > 0:
             self.export()
         if self.report:
             self.generate_report()
@@ -319,7 +336,7 @@ def __init__(
         self.config = config
         self.model_data: Dict[str, Any]
         self.model_path = model_path
-        self._deployment_model: RecursiveScriptModule
+        self._deployment_model: Union[RecursiveScriptModule, Module]
         self.deployment_model_type: str
         self.model_info_filename = "model.json"
         self.report_path = ""
@@ -333,9 +350,7 @@ def deployment_model(self) -> Union[RecursiveScriptModule, nn.Module]:
     @deployment_model.setter
     def deployment_model(self, model_path: str):
         """Set the deployment model."""
-        self._deployment_model, self.deployment_model_type = import_deployment_model(  # type: ignore[assignment]
-            model_path, self.device
-        )
+        self._deployment_model, self.deployment_model_type = import_deployment_model(model_path, self.device)
 
     def prepare(self) -> None:
         """Prepare the evaluation."""
@@ -345,18 +360,24 @@ def prepare(self) -> None:
         if not isinstance(self.model_data, dict):
             raise ValueError("Model info file is not a valid json")
 
-        if self.model_data["input_size"][0] != self.config.transforms.input_width:
-            log.warning(
-                f"Input width of the model ({self.model_data['input_size'][0]}) is different from the one specified "
-                + f"in the config ({self.config.transforms.input_width}). Fixing the config."
-            )
-            self.config.transforms.input_width = self.model_data["input_size"][0]
-
-        if self.model_data["input_size"][1] != self.config.transforms.input_height:
-            log.warning(
-                f"Input height of the model ({self.model_data['input_size'][1]}) is different from the one specified "
-                + f"in the config ({self.config.transforms.input_height}). Fixing the config."
-            )
-            self.config.transforms.input_height = self.model_data["input_size"][1]
+        for input_size in self.model_data["input_size"]:
+            if len(input_size) != 3:
+                continue
+
+            # Adjust the transform for 2D models (CxHxW)
+            # We assume that each input size has the same height and width
+            if input_size[1] != self.config.transforms.input_height:
+                log.warning(
+                    f"Input height of the model ({input_size[1]}) is different from the one specified "
+                    + f"in the config ({self.config.transforms.input_height}). Fixing the config."
+                )
+                self.config.transforms.input_height = input_size[1]
+
+            if input_size[2] != self.config.transforms.input_width:
+                log.warning(
+                    f"Input width of the model ({input_size[2]}) is different from the one specified "
+                    + f"in the config ({self.config.transforms.input_width}). Fixing the config."
+                )
+                self.config.transforms.input_width = input_size[2]
 
         self.deployment_model = self.model_path  # type: ignore[assignment]
diff --git a/quadra/tasks/classification.py b/quadra/tasks/classification.py
index 9c18bd00..cf6d7144 100644
--- a/quadra/tasks/classification.py
+++ b/quadra/tasks/classification.py
@@ -29,6 +29,7 @@
     SklearnClassificationDataModule,
 )
 from quadra.datasets.classification import ImageClassificationListDataset
+from quadra.models.base import ModelSignatureWrapper
 from quadra.models.classification import BaseNetworkBuilder
 from quadra.modules.classification import ClassificationModule
 from quadra.tasks.base import Evaluation, LightningTask, Task
@@ -58,7 +59,11 @@ class Classification(Generic[ClassificationDataModuleT], LightningTask[Classific
     Args:
         config: The experiment configuration
         output: The otuput configuration.
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
         gradcam: Whether to compute gradcams
         checkpoint_path: The path to the checkpoint to load the model from. Defaults to None.
         lr_multiplier: The multiplier for the backbone learning rate. Defaults to None.
@@ -74,7 +79,7 @@ def __init__(
         output: DictConfig,
         checkpoint_path: Optional[str] = None,
         lr_multiplier: Optional[float] = None,
-        export_type: Optional[List[str]] = None,
+        export_config: Optional[DictConfig] = None,
         gradcam: bool = False,
         report: bool = False,
         run_test: bool = False,
@@ -84,7 +89,7 @@ def __init__(
             checkpoint_path=checkpoint_path,
             run_test=run_test,
             report=report,
-            export_type=export_type,
+            export_config=export_config,
         )
         self.output = output
         self.gradcam = gradcam
@@ -244,11 +249,9 @@ def test(self) -> None:
 
     def export(self) -> None:
         """Generate deployment models for the task."""
-        if self.export_type is None:
-            raise ValueError("`Export_type` must be specified in the config to export the model")
-
-        input_width = self.config.transforms.get("input_width")
-        input_height = self.config.transforms.get("input_height")
+        if self.export_config is None or len(self.export_config.types) == 0:
+            log.info("No export type specified skipping export")
+            return
 
         if self.datamodule.class_to_idx is None:
             log.warning(
@@ -257,15 +260,10 @@ def export(self) -> None:
             idx_to_class = {}
         else:
             idx_to_class = {v: k for k, v in self.datamodule.class_to_idx.items()}
-        self.model_json = {
-            "input_size": [input_width, input_height, 3],
-            "classes": idx_to_class,
-            "mean": list(self.config.transforms.mean),
-            "std": list(self.config.transforms.std),
-        }
 
         if self.trainer.checkpoint_callback is None:
             raise ValueError("No checkpoint callback found in the trainer")
+
         best_model_path = self.trainer.checkpoint_callback.best_model_path  # type: ignore[attr-defined]
         log.info("Saving deployment model for %s checkpoint", best_model_path)
 
@@ -278,14 +276,25 @@ def export(self) -> None:
             gradcam=False,
         )
 
-        for export_type in self.export_type:
+        input_shapes = self.export_config.input_shapes
+
+        # TODO: This breaks with bf16 precision!!!
+        half_precision = int(self.trainer.precision) == 16
+
+        for export_type in self.export_config.types:
             if export_type == "torchscript":
-                export_torchscript_model(
-                    module.model,
-                    (1, 3, input_height, input_width),
-                    self.export_folder,
-                    half_precision=int(self.trainer.precision) == 16,
+                out = export_torchscript_model(
+                    model=module.model,
+                    input_shapes=input_shapes,
+                    output_path=self.export_folder,
+                    half_precision=half_precision,
                 )
+
+                if out is None:
+                    log.warning("Skipping torchscript export since the model is not supported")
+                    continue
+
+                _, input_shapes = out
             elif export_type == "pytorch":
                 export_pytorch_model(
                     model=module.model,
@@ -296,6 +305,13 @@ def export(self) -> None:
             else:
                 log.warning("Export type: %s not implemented", export_type)
 
+        self.model_json = {
+            "input_size": input_shapes,
+            "classes": idx_to_class,
+            "mean": list(self.config.transforms.mean),
+            "std": list(self.config.transforms.std),
+        }
+
         with open(os.path.join(self.export_folder, self.deploy_info_file), "w") as f:
             json.dump(self.model_json, f)
 
@@ -426,7 +442,11 @@ class SklearnClassification(Generic[SklearnClassificationDataModuleT], Task[Skle
         config: The experiment configuration
         device: The device to use. Defaults to None.
         output: Dictionary defining which kind of outputs to generate. Defaults to None.
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
@@ -434,9 +454,9 @@ def __init__(
         config: DictConfig,
         output: DictConfig,
         device: str,
-        export_type: Optional[List[str]] = None,
+        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_type=export_type)
+        super().__init__(config=config, export_config=export_config)
 
         self._device = device
         self.output = output
@@ -498,8 +518,10 @@ def backbone(self, backbone_config):
         else:
             log.info("Loading backbone from <%s>", backbone_config.model["_target_"])
             self._backbone = hydra.utils.instantiate(backbone_config.model)
+
+        self._backbone = ModelSignatureWrapper(self._backbone)
         self._backbone.eval()
-        self._backbone = self._backbone.to(self.device)
+        self._backbone.to(self.device)
 
     @property
     def trainer(self) -> SklearnClassificationTrainer:
@@ -634,18 +656,26 @@ def test_full_data(self) -> None:
 
     def export(self) -> None:
         """Generate deployment model for the task."""
-        input_width = self.config.transforms.get("input_width")
-        input_height = self.config.transforms.get("input_height")
-
-        if self.export_type is None:
-            log.info("No export type specified, skipping export")
+        if self.export_config is None or len(self.export_config.types) == 0:
+            log.info("No export type specified skipping export")
             return
 
-        for export_type in self.export_type:
+        input_shapes = self.export_config.input_shapes
+
+        for export_type in self.export_config.types:
             if export_type == "torchscript":
-                export_torchscript_model(
-                    self.backbone, (1, 3, input_height, input_width), self.export_folder, half_precision=False
+                out = export_torchscript_model(
+                    model=self.backbone,
+                    input_shapes=input_shapes,
+                    output_path=self.export_folder,
+                    half_precision=False,
                 )
+
+                if out is None:
+                    log.warning("Skipping torchscript export since the model is not supported")
+                    continue
+
+                _, input_shapes = out
             elif export_type == "pytorch":
                 os.makedirs(self.export_folder, exist_ok=True)
                 # We only need to save classifier.joblib + backbone config file
@@ -660,7 +690,7 @@ def export(self) -> None:
         idx_to_class = {v: k for k, v in self.datamodule.full_dataset.class_to_idx.items()}
 
         model_json = {
-            "input_size": [input_width, input_height, 3],
+            "input_size": input_shapes,
             "classes": idx_to_class,
             "mean": list(self.config.transforms.mean),
             "std": list(self.config.transforms.std),
@@ -712,7 +742,7 @@ def execute(self) -> None:
         if self.output.report:
             self.generate_report()
         self.train_full_data()
-        if self.export_type is not None and len(self.export_type) > 0:
+        if self.export_config is not None and len(self.export_config.types) > 0:
             self.export()
         if self.output.test_full_data:
             self.test_full_data()
@@ -966,7 +996,7 @@ def deployment_model(self, model_path: str):
         if os.path.splitext(os.path.basename(model_path))[1] == ".pth":
             model_config = OmegaConf.load(os.path.join(Path(self.model_path).parent, "model_config.yaml"))
             self.model = model_config  # type: ignore[assignment]
-        self._deployment_model, self.deployment_model_type = import_deployment_model(  # type: ignore[assignment]
+        self._deployment_model, self.deployment_model_type = import_deployment_model(
             model_path, self.device, self.model
         )
         if self.gradcam and self.deployment_model_type != "torch":
@@ -1036,6 +1066,7 @@ def test(self) -> None:
                             zoom_factors = tuple(np.array(new_shape) / np.array(orig_shape))
                             grayscale_cam = ndimage.zoom(grayscale_cam_low_res, zoom_factors, order=1)
                             grayscale_cams_list.append(torch.from_numpy(grayscale_cam))
+
         grayscale_cams: Optional[torch.Tensor] = None
         if self.gradcam:
             grayscale_cams = torch.cat(grayscale_cams_list, dim=0)
diff --git a/quadra/tasks/patch.py b/quadra/tasks/patch.py
index 9d6bec1b..afb0def6 100644
--- a/quadra/tasks/patch.py
+++ b/quadra/tasks/patch.py
@@ -11,6 +11,7 @@
 
 from quadra.datamodules import PatchSklearnClassificationDataModule
 from quadra.datasets.patch import PatchSklearnClassificationTrainDataset
+from quadra.models.base import ModelSignatureWrapper
 from quadra.tasks.base import Evaluation, Task
 from quadra.trainers.classification import SklearnClassificationTrainer
 from quadra.utils import utils
@@ -28,7 +29,11 @@ class PatchSklearnClassification(Task[PatchSklearnClassificationDataModule]):
         config: The experiment configuration
         device: The device to use
         output: Dictionary defining which kind of outputs to generate. Defaults to None.
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
@@ -36,9 +41,9 @@ def __init__(
         config: DictConfig,
         output: DictConfig,
         device: str,
-        export_type: Optional[List[str]] = None,
+        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_type=export_type)
+        super().__init__(config=config, export_config=export_config)
         self.device: str = device
         self.output: DictConfig = output
         self.return_polygon: bool = True
@@ -79,6 +84,8 @@ def backbone(self, backbone_config):
         else:
             log.info("Loading backbone from <%s>", backbone_config.model["_target_"])
             self._backbone = hydra.utils.instantiate(backbone_config.model)
+
+        self._backbone = ModelSignatureWrapper(self._backbone)
         self._backbone.eval()
         self._backbone = self._backbone.to(self.device)
 
@@ -200,18 +207,29 @@ def generate_report(self) -> None:
 
     def export(self) -> None:
         """Generate deployment model for the task."""
-        if self.export_type is None or len(self.export_type) == 0:
+        if self.export_config is None or len(self.export_config.types) == 0:
             log.info("No export type specified skipping export")
             return
-        input_width = self.config.transforms.get("input_width")
-        input_height = self.config.transforms.get("input_height")
 
-        for export_type in self.export_type:
+        os.makedirs(self.export_folder, exist_ok=True)
+
+        input_shapes = self.export_config.input_shapes
+
+        for export_type in self.export_config.types:
             if export_type == "torchscript":
-                export_torchscript_model(
-                    self.backbone, (1, 3, input_height, input_width), self.export_folder, half_precision=False
+                out = export_torchscript_model(
+                    model=self.backbone,
+                    input_shapes=input_shapes,
+                    output_path=self.export_folder,
+                    half_precision=False,
                 )
 
+                if out is None:
+                    log.warning("Skipping torchscript export since the model is not supported")
+                    continue
+
+                _, input_shapes = out
+
         dump(self.model, os.path.join(self.export_folder, "classifier.joblib"))
 
         dataset_info = self.datamodule.info
@@ -225,7 +243,7 @@ def export(self) -> None:
         idx_to_class = {v: k for k, v in self.datamodule.class_to_idx.items()}
 
         model_json = {
-            "input_size": [input_width, input_height, 3],
+            "input_size": input_shapes,
             "classes": idx_to_class,
             "mean": self.config.transforms.mean,
             "std": self.config.transforms.std,
@@ -247,7 +265,7 @@ def execute(self) -> None:
         self.train()
         if self.output.report:
             self.generate_report()
-        if self.export_type is not None and len(self.export_type) > 0:
+        if self.export_config is not None and len(self.export_config.types) > 0:
             self.export()
         self.finalize()
 
diff --git a/quadra/tasks/segmentation.py b/quadra/tasks/segmentation.py
index c3d977ae..11de316c 100644
--- a/quadra/tasks/segmentation.py
+++ b/quadra/tasks/segmentation.py
@@ -11,6 +11,7 @@
 
 from quadra.callbacks.mlflow import get_mlflow_logger
 from quadra.datamodules import SegmentationDataModule, SegmentationMulticlassDataModule
+from quadra.models.base import ModelSignatureWrapper
 from quadra.modules.base import SegmentationModel
 from quadra.tasks.base import Evaluation, LightningTask
 from quadra.utils import utils
@@ -34,7 +35,11 @@ class Segmentation(Generic[SegmentationDataModuleT], LightningTask[SegmentationD
         run_test: If True, run test after training. Defaults to False.
         evaluate: Dict with evaluation parameters. Defaults to None.
         report: If True, create report after training. Defaults to False.
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
@@ -45,29 +50,34 @@ def __init__(
         run_test: bool = False,
         evaluate: Optional[DictConfig] = None,
         report: bool = False,
-        export_type: Optional[List[str]] = None,
+        export_config: Optional[DictConfig] = None,
     ):
         super().__init__(
             config=config,
             checkpoint_path=checkpoint_path,
             run_test=run_test,
             report=report,
-            export_type=export_type,
+            export_config=export_config,
         )
         self.evaluate = evaluate
         self.num_viz_samples = num_viz_samples
         self.export_folder: str = "deployment_model"
         self.exported_model_path: Optional[str] = None
         if self.evaluate and any(self.evaluate.values()):
-            if self.export_type is None or len(self.export_type) == 0 or "torchscript" not in self.export_type:
+            if (
+                self.export_config is None
+                or len(self.export_config.types) == 0
+                or "torchscript" not in self.export_config.types
+            ):
                 log.info(
                     "Evaluation is enabled, but training does not export a deployment model. Automatically export the "
                     "model as torchscript."
                 )
-                if self.export_type is None:
-                    self.export_type = ["torchscript"]
+                if self.export_config is None:
+                    self.export_config = DictConfig({"types": ["torchscript"]})
                 else:
-                    self.export_type.append("torchscript")
+                    self.export_config.types.append("torchscript")
+
             if not self.report:
                 log.info("Evaluation is enabled, but reporting is disabled. Enabling reporting automatically.")
                 self.report = True
@@ -91,6 +101,7 @@ def module(self, module_config) -> None:
                 module_config.model.num_classes = len(self.datamodule.idx_to_class) + 1
 
         model = hydra.utils.instantiate(module_config.model)
+        model = ModelSignatureWrapper(model)
         log.info("Instantiating optimizer <%s>", self.config.optimizer["_target_"])
         param_list = []
         for param in model.parameters():
@@ -113,9 +124,6 @@ def prepare(self) -> None:
     def export(self) -> None:
         """Generate a deployment model for the task."""
         log.info("Exporting model ready for deployment")
-        input_width = self.config.transforms.get("input_width")
-        input_height = self.config.transforms.get("input_height")
-
         # Get best model!
         if self.trainer.checkpoint_callback is None:
             raise ValueError("No checkpoint callback found in the trainer")
@@ -136,19 +144,37 @@ def export(self) -> None:
         else:
             log.info("idx_to_class is present")
             classes = self.config.datamodule.idx_to_class
-        if self.export_type is None:
+
+        if self.export_config is None:
             raise ValueError(
                 "No export type specified. This should not happen, please check if you have set "
                 "the export_type or assign it to a default value."
             )
-        for export_type in self.export_type:
+
+        input_shapes = self.export_config.input_shapes
+
+        half_precision = self.trainer.precision == 16
+
+        for export_type in self.export_config.types:
             if export_type == "torchscript":
-                self.exported_model_path = export_torchscript_model(
-                    module.model, (1, 3, input_height, input_width), self.export_folder, half_precision=False
+                out = export_torchscript_model(
+                    model=module.model,
+                    input_shapes=input_shapes,
+                    output_path=self.export_folder,
+                    half_precision=half_precision,
                 )
 
+                if out is None:
+                    log.warning("Skipping torchscript export since the model is not supported")
+                    continue
+
+                self.exported_model_path, input_shapes = out
+
+        if input_shapes is None:
+            log.warning("Not able to export the model in any format")
+
         model_json = {
-            "input_size": [input_width, input_height, 3],
+            "input_size": input_shapes,
             "classes": classes,
             "mean": self.config.transforms.mean,
             "std": self.config.transforms.std,
@@ -237,8 +263,9 @@ def prepare(self) -> None:
         self.config.transforms.mean = self.model_data["mean"]
         self.config.transforms.std = self.model_data["std"]
         # Setup datamodule
-        idx_to_class = self.model_data["classes"]  # dict {index: class}
-        self.config.datamodule.idx_to_class = idx_to_class
+        if hasattr(self.config.datamodule, "idx_to_class"):
+            idx_to_class = self.model_data["classes"]  # dict {index: class}
+            self.config.datamodule.idx_to_class = idx_to_class
         self.datamodule = self.config.datamodule
         # prepare_data() must be explicitly called because there is no lightning training
         self.datamodule.prepare_data()
diff --git a/quadra/tasks/ssl.py b/quadra/tasks/ssl.py
index 290e8dd2..84cd3af8 100644
--- a/quadra/tasks/ssl.py
+++ b/quadra/tasks/ssl.py
@@ -1,6 +1,6 @@
 import json
 import os
-from typing import Any, List, Optional, Tuple, cast
+from typing import Any, List, Optional, Tuple, Union, cast
 
 import hydra
 import torch
@@ -13,6 +13,7 @@
 from tqdm import tqdm
 
 from quadra.callbacks.scheduler import WarmupInit
+from quadra.models.base import ModelSignatureWrapper
 from quadra.tasks.base import LightningTask, Task
 from quadra.utils import utils
 from quadra.utils.export import export_torchscript_model, import_deployment_model
@@ -28,7 +29,11 @@ class SSL(LightningTask):
         checkpoint_path: The path to the checkpoint to load the model from Defaults to None
         report: Whether to create the report
         run_test: Whether to run final test
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
@@ -37,14 +42,14 @@ def __init__(
         run_test: bool = False,
         report: bool = False,
         checkpoint_path: Optional[str] = None,
-        export_type: Optional[List[str]] = None,
+        export_config: Optional[DictConfig] = None,
     ):
         super().__init__(
             config=config,
             checkpoint_path=checkpoint_path,
             run_test=run_test,
             report=report,
-            export_type=export_type,
+            export_config=export_config,
         )
         self._backbone: nn.Module
         self._optimizer: torch.optim.Optimizer
@@ -95,29 +100,33 @@ def test(self) -> None:
 
     def export(self) -> None:
         """Deploy a model ready for production."""
-        if self.export_type is None or len(self.export_type) == 0:
+        if self.export_config is None or len(self.export_config.types) == 0:
             log.info("No export type specified skipping export")
             return
 
-        if self.export_type is None:
-            return
-
-        input_width = self.config.transforms.get("input_width")
-        input_height = self.config.transforms.get("input_height")
         mean = self.config.transforms.mean
         std = self.config.transforms.std
-        half_precision = self.trainer.precision == 16
-
-        if "torchscript" in self.export_type:
-            export_torchscript_model(
-                model=cast(nn.Module, self.module.model),
-                input_shape=(1, 3, input_height, input_width),
-                half_precision=half_precision,
-                output_path=self.export_folder,
-            )
+        half_precision = int(self.trainer.precision) == 16
+
+        input_shapes = self.export_config.input_shapes
+
+        for export_type in self.export_config.types:
+            if export_type == "torchscript":
+                out = export_torchscript_model(
+                    model=cast(nn.Module, self.module.model),
+                    input_shapes=input_shapes,
+                    output_path=self.export_folder,
+                    half_precision=half_precision,
+                )
+
+                if out is None:
+                    log.warning("Skipping torchscript export since the model is not supported")
+                    continue
+
+                _, input_shapes = out
 
         model_json = {
-            "input_size": [3, input_width, input_height],
+            "input_size": input_shapes,
             "classes": None,
             "mean": mean,
             "std": std,
@@ -132,21 +141,25 @@ class Simsiam(SSL):
 
     Args:
         config: the main config
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
         checkpoint_path: if a checkpoint is specified, then it will return a trained model,
             with weights loaded from the checkpoint path specified.
             Defaults to None.
         run_test: Whether to run final test
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
         self,
         config: DictConfig,
-        export_type: Optional[List[str]] = None,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
+        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_type=export_type, checkpoint_path=checkpoint_path, run_test=run_test)
+        super().__init__(config=config, export_config=export_config, checkpoint_path=checkpoint_path, run_test=run_test)
         self.backbone: nn.Module
         self.projection_mlp: nn.Module
         self.prediction_mlp: nn.Module
@@ -203,21 +216,25 @@ class SimCLR(SSL):
 
     Args:
         config: the main config
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
         checkpoint_path: if a checkpoint is specified, then it will return a trained model,
             with weights loaded from the checkpoint path specified.
             Defaults to None.
         run_test: Whether to run final test
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
         self,
         config: DictConfig,
-        export_type: Optional[List[str]] = None,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
+        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_type=export_type, checkpoint_path=checkpoint_path, run_test=run_test)
+        super().__init__(config=config, export_config=export_config, checkpoint_path=checkpoint_path, run_test=run_test)
         self.backbone: nn.Module
         self.projection_mlp: nn.Module
 
@@ -258,6 +275,7 @@ def module(self, module_config):
                 lr_scheduler=self.scheduler,
             )
         self._module = module
+        self._module.model = ModelSignatureWrapper(self._module.model)
 
 
 class Barlow(SimCLR):
@@ -265,22 +283,25 @@ class Barlow(SimCLR):
 
     Args:
         config: the main config
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
         checkpoint_path: if a checkpoint is specified, then it will return a trained model,
             with weights loaded from the checkpoint path specified.
             Defaults to None.
         run_test: Whether to run final test
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
 
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
         self,
         config: DictConfig,
-        export_type: Optional[List[str]] = None,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
+        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_type=export_type, checkpoint_path=checkpoint_path, run_test=run_test)
+        super().__init__(config=config, export_config=export_config, checkpoint_path=checkpoint_path, run_test=run_test)
 
     def prepare(self) -> None:
         """Prepare the experiment."""
@@ -305,26 +326,29 @@ class BYOL(SSL):
 
     Args:
         config: the main config
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
         checkpoint_path: if a checkpoint is specified, then it will return a trained model,
             with weights loaded from the checkpoint path specified.
             Defaults to None.
         run_test: Whether to run final test
-        **kwargs: Keyword arguments
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
 
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
+        **kwargs: Keyword arguments
     """
 
     def __init__(
         self,
         config: DictConfig,
-        export_type: Optional[List[str]] = None,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
+        export_config: Optional[DictConfig] = None,
         **kwargs: Any,
     ):
         super().__init__(
             config=config,
-            export_type=export_type,
+            export_config=export_config,
             checkpoint_path=checkpoint_path,
             run_test=run_test,
             **kwargs,
@@ -392,21 +416,25 @@ class DINO(SSL):
 
     Args:
         config: the main config
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
         checkpoint_path: if a checkpoint is specified, then it will return a trained model,
             with weights loaded from the checkpoint path specified.
             Defaults to None.
         run_test: Whether to run final test
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
         self,
         config: DictConfig,
-        export_type: Optional[List[str]] = None,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
+        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_type=export_type, checkpoint_path=checkpoint_path, run_test=run_test)
+        super().__init__(config=config, export_config=export_config, checkpoint_path=checkpoint_path, run_test=run_test)
         self.student_model: nn.Module
         self.teacher_model: nn.Module
         self.student_projection_mlp: nn.Module
@@ -501,7 +529,7 @@ def __init__(
         self.embedding_writer = SummaryWriter(self.embeddings_path)
         self.writer_step = 0  # for tensorboard
         self.embedding_image_size = embedding_image_size
-        self._deployment_model: RecursiveScriptModule
+        self._deployment_model: Union[RecursiveScriptModule, nn.Module]
         self.deployment_model_type: str
 
     @property
diff --git a/quadra/utils/export.py b/quadra/utils/export.py
index 3c4222ab..c66040c4 100644
--- a/quadra/utils/export.py
+++ b/quadra/utils/export.py
@@ -1,50 +1,93 @@
 import os
-from typing import Optional, Tuple, Union, cast
+from typing import Any, List, Optional, Sequence, Tuple, Union, cast
 
 import torch
 from anomalib.models.cflow import CflowLightning
 from torch import nn
 from torch.jit._script import RecursiveScriptModule
 
-from quadra.utils.utils import get_logger
+from quadra.models.base import ModelSignatureWrapper
 
-log = get_logger(__name__)
+# TODO: Solve circular import as it is not possible to import get_logger right now
+
+
+def generate_torch_inputs(
+    input_shapes: List[Any], device: str, half_precision: bool = False, dtype: torch.dtype = torch.float32
+) -> Union[List[Any], Tuple[Any, ...], torch.Tensor]:
+    """Given a list of input shapes that can contain either lists, tuples or dicts, with tuples being the input shapes
+    of the model, generate a list of torch tensors with the given device and dtype.
+    """
+    if isinstance(input_shapes, list):
+        if any(isinstance(inp, (Sequence, dict)) for inp in input_shapes):
+            return [generate_torch_inputs(inp, device, half_precision, dtype) for inp in input_shapes]
+
+        # Base case
+        inp = torch.randn((1, *input_shapes), dtype=dtype, device=device)
+
+    if isinstance(input_shapes, dict):
+        return {k: generate_torch_inputs(v, device, half_precision, dtype) for k, v in input_shapes.items()}
+
+    if isinstance(input_shapes, tuple):
+        if any(isinstance(inp, (Sequence, dict)) for inp in input_shapes):
+            # The tuple contains a list, tuple or dict
+            return tuple(generate_torch_inputs(inp, device, half_precision, dtype) for inp in input_shapes)
+
+        # Base case
+        inp = torch.randn((1, *input_shapes), dtype=dtype, device=device)
+
+    if half_precision:
+        inp = inp.half()
+
+    return inp
 
 
 def export_torchscript_model(
     model: nn.Module,
-    input_shape: Tuple[int, ...],
     output_path: str,
+    input_shapes: Optional[List[Any]] = None,
     half_precision: bool = False,
     model_name: str = "model.pt",
-) -> Optional[str]:
+) -> Optional[Tuple[str, Any]]:
     """Export a PyTorch model with TorchScript.
 
     Args:
         model: PyTorch model to be exported
-        input_shape: Input shape for tracing
+        input_shapes: Inputs shape for tracing
         output_path: Path to save the model
         half_precision: If True, the model will be exported with half precision
         model_name: Name of the exported model
 
     Returns:
-        If the model is exported successfully, the path to the model is returned.
+        If the model is exported successfully, the path to the model and the input shape are returned.
 
     """
-    model.eval()
+    if isinstance(model, ModelSignatureWrapper):
+        if input_shapes is None:
+            input_shapes = model.input_shapes
+        model = model.instance
+
+    if input_shapes is None:
+        # log.warning(
+        #    "Input shape is None, can not trace model! Please provide input_shapes in the task export configuration."
+        # )
+        return None
+
     if isinstance(model, CflowLightning):
-        log.warning("Exporting cflow model with torchscript is not supported yet.")
+        # log.warning("Exporting cflow model with torchscript is not supported yet.")
         return None
 
+    model.eval()
     if half_precision:
-        log.info("Jitting model with half precision on GPU")
+        # log.info("Jitting model with half precision on GPU")
         model.to("cuda:0")
         model = model.half()
-        inp = torch.randn(input_shape, dtype=torch.float16, device="cuda:0")
+        inp = generate_torch_inputs(
+            input_shapes=input_shapes, device="cuda:0", half_precision=True, dtype=torch.float16
+        )
     else:
-        log.info("Jitting model with double precision")
+        # log.info("Jitting model with double precision")
         model.cpu()
-        inp = torch.randn(input_shape, dtype=torch.float32)
+        inp = generate_torch_inputs(input_shapes=input_shapes, device="cpu", half_precision=False, dtype=torch.float32)
 
     with torch.no_grad():
         model_jit = torch.jit.trace(model, inp)
@@ -54,9 +97,9 @@ def export_torchscript_model(
     model_path = os.path.join(output_path, model_name)
     model_jit.save(model_path)
 
-    log.info("Torchscript model saved to %s", os.path.join(os.getcwd(), model_path))
+    # log.info("Torchscript model saved to %s", os.path.join(os.getcwd(), model_path))
 
-    return os.path.join(os.getcwd(), model_path)
+    return os.path.join(os.getcwd(), model_path), input_shapes
 
 
 def export_pytorch_model(model: nn.Module, output_path: str, model_name: str = "model.pth") -> str:
@@ -76,7 +119,7 @@ def export_pytorch_model(model: nn.Module, output_path: str, model_name: str = "
     model.cpu()
     model_path = os.path.join(output_path, model_name)
     torch.save(model.state_dict(), model_path)
-    log.info("Pytorch model saved to %s", os.path.join(output_path, model_name))
+    # log.info("Pytorch model saved to %s", os.path.join(output_path, model_name))
 
     return os.path.join(os.getcwd(), model_path)
 
@@ -96,6 +139,7 @@ def import_deployment_model(
     Returns:
         A tuple containing the model and the model type
     """
+
     file_extension = os.path.splitext(os.path.basename(model_path))[1]
     if file_extension == ".pt":
         model = cast(RecursiveScriptModule, torch.jit.load(model_path))
@@ -105,12 +149,12 @@ def import_deployment_model(
         return model, "torchscript"
     if file_extension == ".pth":
         if model is None:
-            log.warning("Model is None, can not load state_dict")
-        else:
-            model.load_state_dict(torch.load(model_path))
-            model.eval()
-            model.to(device)
+            raise ValueError("Model is not defined, can not load state_dict!")
+
+        model.load_state_dict(torch.load(model_path))
+        model.eval()
+        model.to(device)
 
-            return model, "torch"
+        return model, "torch"
 
     raise ValueError(f"Unable to load model with extension {file_extension}, valid extensions are: ['.pt', 'pth']")
diff --git a/quadra/utils/mlflow.py b/quadra/utils/mlflow.py
new file mode 100644
index 00000000..438e1f1e
--- /dev/null
+++ b/quadra/utils/mlflow.py
@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+try:
+    from mlflow.models import infer_signature  # noqa
+    from mlflow.models.signature import ModelSignature  # noqa
+
+    MLFLOW_AVAILABLE = True
+except ImportError:
+    MLFLOW_AVAILABLE = False
+
+from typing import Any, Sequence, TypeVar
+
+import torch
+from pytorch_lightning import Trainer
+from pytorch_lightning.loggers import MLFlowLogger
+from torch import nn
+
+NnModuleT = TypeVar("NnModuleT", bound=nn.Module)
+
+
+@torch.inference_mode()
+def infer_signature_torch_model(model: NnModuleT, data: list[Any]) -> ModelSignature | None:
+    """Infer input and output signature for a PyTorch/Torchscript model."""
+    model = model.eval()
+    model = model.cpu()
+    model_output = model(*data)
+
+    try:
+        output_signature = infer_signature_input_torch(model_output)
+
+        if len(data) == 1:
+            signature_input = infer_signature_input_torch(data[0])
+        else:
+            signature_input = infer_signature_input_torch(data)
+    except ValueError:
+        # TODO: Solve circular import as it is not possible to import get_logger right now
+        # log.warning("Unable to infer signature for model output type %s", type(model_output))
+        return None
+
+    return infer_signature(signature_input, output_signature)
+
+
+def infer_signature_input_torch(input_tensor: Any) -> Any:
+    """Recursively infer the signature input format to pass to mlflow.models.infer_signature.
+
+    Raises:
+        ValueError: If the input type is not supported or when nested dicts or sequences are encountered.
+    """
+    if isinstance(input_tensor, Sequence):
+        # Mlflow currently does not support sequence outputs, so we use a dict instead
+        signature = {}
+        for i, x in enumerate(input_tensor):
+            if isinstance(x, Sequence):
+                # Nested signature is currently not supported by mlflow
+                raise ValueError("Nested sequences are not supported")
+                # TODO: Enable this once mlflow supports nested signatures
+                # signature[f"output_{i}"] = {f"output_{j}": infer_signature_torch(y) for j, y in enumerate(x)}
+            if isinstance(x, dict):
+                # Nested dicts are not supported
+                raise ValueError("Nested dicts are not supported")
+
+            signature[f"output_{i}"] = infer_signature_input_torch(x)
+    elif isinstance(input_tensor, torch.Tensor):
+        signature = input_tensor.cpu().numpy()
+    elif isinstance(input_tensor, dict):
+        signature = {}
+        for k, v in input_tensor.items():
+            if isinstance(v, dict):
+                # Nested dicts are not supported
+                raise ValueError("Nested dicts are not supported")
+            if isinstance(v, Sequence):
+                # Nested signature is currently not supported by mlflow
+                raise ValueError("Nested sequences are not supported")
+
+            signature[k] = infer_signature_input_torch(v)
+    else:
+        raise ValueError(f"Unable to infer signature for model output type {type(input_tensor)}")
+
+    return signature
+
+
+def get_mlflow_logger(trainer: Trainer) -> MLFlowLogger | None:
+    """Safely get Mlflow logger from Trainer loggers.
+
+    Args:
+        trainer: Pytorch Lightning trainer.
+
+    Returns:
+        An mlflow logger if available, else None.
+    """
+    if isinstance(trainer.logger, MLFlowLogger):
+        return trainer.logger
+
+    if isinstance(trainer.logger, list):
+        for logger in trainer.logger:
+            if isinstance(logger, MLFlowLogger):
+                return logger
+
+    return None
diff --git a/quadra/utils/model_manager.py b/quadra/utils/model_manager.py
new file mode 100644
index 00000000..1dbfce1f
--- /dev/null
+++ b/quadra/utils/model_manager.py
@@ -0,0 +1,321 @@
+from __future__ import annotations
+
+import getpass
+import os
+from abc import ABC, abstractmethod
+from datetime import datetime
+from typing import Any, Literal
+
+from quadra.utils.utils import get_logger
+
+log = get_logger(__name__)
+
+try:
+    import mlflow  # noqa
+    from mlflow.entities import Run  # noqa
+    from mlflow.entities.model_registry import ModelVersion  # noqa
+    from mlflow.exceptions import RestException  # noqa
+    from mlflow.tracking import MlflowClient  # noqa
+
+    MLFLOW_AVAILABLE = True
+except ImportError:
+    MLFLOW_AVAILABLE = False
+
+
+VERSION_MD_TEMPLATE = "## **Version {}**\n"
+DESCRIPTION_MD_TEMPLATE = "### Description: \n{}\n"
+
+
+class AbstractModelManager(ABC):
+    """Abstract class for model managers."""
+
+    @abstractmethod
+    def register_model(
+        self, model_location: str, model_name: str, description: str, tags: dict[str, Any] | None = None
+    ) -> Any:
+        """Register a model in the model registry."""
+
+    @abstractmethod
+    def get_latest_version(self, model_name: str) -> Any:
+        """Get the latest version of a model for all the possible stages or filtered by stage."""
+
+    @abstractmethod
+    def transition_model(self, model_name: str, version: int, stage: str, description: str | None = None) -> Any:
+        """Transition the model with the given version to a new stage."""
+
+    @abstractmethod
+    def delete_model(self, model_name: str, version: int, description: str | None = None) -> None:
+        """Delete a model with the given version."""
+
+    @abstractmethod
+    def register_best_model(
+        self,
+        experiment_name: str,
+        metric: str,
+        model_name: str,
+        description: str,
+        tags: dict[str, Any] | None = None,
+        mode: Literal["max", "min"] = "max",
+        model_path: str = "deployment_model",
+    ) -> Any:
+        """Register the best model from an experiment."""
+
+    @abstractmethod
+    def download_model(self, model_name: str, version: int, output_path: str) -> None:
+        """Download the model with the given version to the given output path."""
+
+
+class MlflowModelManager(AbstractModelManager):
+    """Model manager for Mlflow."""
+
+    def __init__(self):
+        if not MLFLOW_AVAILABLE:
+            raise ImportError("Mlflow is not available, please install it with pip install mlflow")
+
+        if os.getenv("MLFLOW_TRACKING_URI") is None:
+            raise ValueError("MLFLOW_TRACKING_URI environment variable is not set")
+
+        mlflow.set_tracking_uri(os.getenv("MLFLOW_TRACKING_URI"))
+        self.client = MlflowClient()
+
+    def register_model(
+        self, model_location: str, model_name: str, description: str | None = None, tags: dict[str, Any] | None = None
+    ) -> ModelVersion:
+        """Register a model in the model registry.
+
+        Args:
+            model_location: The model uri
+            model_name: The name of the model after it is registered
+            description: A description of the model, this will be added to the model changelog
+            tags: A dictionary of tags to add to the model
+
+        Returns:
+            The model version
+        """
+        model_version = mlflow.register_model(model_uri=model_location, name=model_name, tags=tags)
+        log.info("Registered model %s with version %s", model_name, model_version.version)
+        registered_model_description = self.client.get_registered_model(model_name).description
+
+        if model_version.version == "1":
+            header = "# MODEL CHANGELOG\n"
+        else:
+            header = ""
+
+        new_model_description = VERSION_MD_TEMPLATE.format(model_version.version)
+        new_model_description += self._get_author_and_date()
+        new_model_description += self._generate_description(description)
+
+        self.client.update_registered_model(model_name, header + registered_model_description + new_model_description)
+
+        self.client.update_model_version(
+            model_name, model_version.version, "# MODEL CHANGELOG\n" + new_model_description
+        )
+
+        return model_version
+
+    def get_latest_version(self, model_name: str) -> ModelVersion:
+        """Get the latest version of a model.
+
+        Args:
+            model_name: The name of the model
+
+        Returns:
+            The model version
+        """
+        latest_version = max(int(x.version) for x in self.client.get_latest_versions("manager_model"))
+        model_version = self.client.get_model_version(model_name, latest_version)
+
+        return model_version
+
+    def transition_model(
+        self, model_name: str, version: int, stage: str, description: str | None = None
+    ) -> ModelVersion | None:
+        """Transition a model to a new stage.
+
+        Args:
+            model_name: The name of the model
+            version: The version of the model
+            stage: The stage of the model
+            description: A description of the transition, this will be added to the model changelog
+        """
+        previous_stage = self._safe_get_stage(model_name, version)
+
+        if previous_stage is None:
+            return None
+
+        if previous_stage.lower() == stage.lower():
+            log.warning("Model %s version %s is already in stage %s", model_name, version, stage)
+            return self.client.get_model_version(model_name, version)
+
+        log.info("Transitioning model %s version %s from %s to %s", model_name, version, previous_stage, stage)
+        model_version = self.client.transition_model_version_stage(name=model_name, version=version, stage=stage)
+        new_stage = model_version.current_stage
+        registered_model_description = self.client.get_registered_model(model_name).description
+        single_model_description = self.client.get_model_version(model_name, version).description
+
+        new_model_description = "## **Transition:**\n"
+        new_model_description += f"### Version {model_version.version} from {previous_stage} to {new_stage}\n"
+        new_model_description += self._get_author_and_date()
+        new_model_description += self._generate_description(description)
+
+        self.client.update_registered_model(model_name, registered_model_description + new_model_description)
+        self.client.update_model_version(
+            model_name, model_version.version, single_model_description + new_model_description
+        )
+
+        return model_version
+
+    def delete_model(self, model_name: str, version: int, description: str | None = None) -> None:
+        """Delete a model.
+
+        Args:
+            model_name: The name of the model
+            version: The version of the model
+            description: Why the model was deleted, this will be added to the model changelog
+        """
+        model_stage = self._safe_get_stage(model_name, version)
+
+        if model_stage is None:
+            return
+
+        if (
+            input(
+                f"Model named `{model_name}`, version {version} is in stage {model_stage}, "
+                "type the model name to continue deletion:"
+            )
+            != model_name
+        ):
+            log.warning("Model name did not match, aborting deletion")
+            return
+
+        log.info("Deleting model %s version %s", model_name, version)
+        self.client.delete_model_version(model_name, version)
+
+        registered_model_description = self.client.get_registered_model(model_name).description
+
+        new_model_description = "## **Deletion:**\n"
+        new_model_description += f"### Version {version} from stage: {model_stage}\n"
+        new_model_description += self._get_author_and_date()
+        new_model_description += self._generate_description(description)
+
+        self.client.update_registered_model(model_name, registered_model_description + new_model_description)
+
+    def register_best_model(
+        self,
+        experiment_name: str,
+        metric: str,
+        model_name: str,
+        description: str | None = None,
+        tags: dict[str, Any] | None = None,
+        mode: Literal["max", "min"] = "max",
+        model_path: str = "deployment_model",
+    ) -> ModelVersion | None:
+        """Register the best model from an experiment.
+
+        Args:
+            experiment_name: The name of the experiment
+            metric: The metric to use to determine the best model
+            model_name: The name of the model after it is registered
+            description: A description of the model, this will be added to the model changelog
+            tags: A dictionary of tags to add to the model
+            mode: The mode to use to determine the best model, either "max" or "min"
+            model_path: The path to the model within the experiment run
+
+        Returns:
+            The registered model version if successful, otherwise None
+        """
+        if mode not in ["max", "min"]:
+            raise ValueError(f"Mode must be either 'max' or 'min', got {mode}")
+
+        experiment_id = self.client.get_experiment_by_name(experiment_name).experiment_id
+        runs = self.client.search_runs(experiment_ids=[experiment_id])
+
+        if len(runs) == 0:
+            log.error("No runs found for experiment %s", experiment_name)
+            return None
+
+        best_run: Run | None = None
+
+        # We can only make comparisons if the model is on the top folder, otherwise just check if the folder exists
+        # TODO: Is there a better way to do this?
+        base_model_path = model_path.split("/")[0]
+
+        for run in runs:
+            run_artifacts = [x.path for x in self.client.list_artifacts(run.info.run_id) if x.path == base_model_path]
+
+            if len(run_artifacts) == 0:
+                # If we don't find the given model path, skip this run
+                continue
+
+            if best_run is None:
+                # If we find a run with the model it must also have the metric
+                if run.data.metrics.get(metric) is not None:
+                    best_run = run
+                continue
+
+            if mode == "max":
+                if run.data.metrics[metric] > best_run.data.metrics[metric]:
+                    best_run = run
+            else:
+                if run.data.metrics[metric] < best_run.data.metrics[metric]:
+                    best_run = run
+
+        if best_run is None:
+            log.error("No runs found for experiment %s with the given metric", experiment_name)
+            return None
+
+        best_model_uri = f"runs:/{best_run.info.run_id}/{model_path}"
+
+        model_version = self.register_model(
+            model_location=best_model_uri, model_name=model_name, tags=tags, description=description
+        )
+
+        return model_version
+
+    def download_model(self, model_name: str, version: int, output_path: str) -> None:
+        """Download the model with the given version to the given output path.
+
+        Args:
+            model_name: The name of the model
+            version: The version of the model
+            output_path: The path to save the model to
+        """
+        artifact_uri = self.client.get_model_version_download_uri(model_name, version)
+        log.info("Downloading model %s version %s from %s to %s", model_name, version, artifact_uri, output_path)
+        if not os.path.exists(output_path):
+            log.info("Creating output path %s", output_path)
+            os.makedirs(output_path)
+        mlflow.artifacts.download_artifacts(artifact_uri=artifact_uri, dst_path=output_path)
+
+    @staticmethod
+    def _generate_description(description: str | None = None) -> str:
+        """Generate the description markdown template."""
+        if description is None:
+            return ""
+
+        return DESCRIPTION_MD_TEMPLATE.format(description)
+
+    @staticmethod
+    def _get_author_and_date() -> str:
+        """Get the author and date markdown template."""
+        author_and_date = f"### Author: {getpass.getuser()}\n"
+        author_and_date += f"### Date: {datetime.now().astimezone().strftime('%d/%m/%Y %H:%M:%S %Z')}\n"
+
+        return author_and_date
+
+    def _safe_get_stage(self, model_name: str, version: int) -> str | None:
+        """Get the stage of a model version.
+
+        Args:
+            model_name: The name of the model
+            version: The version of the model
+
+        Returns:
+            The stage of the model version if it exists, otherwise None
+        """
+        try:
+            model_stage = self.client.get_model_version(model_name, version).current_stage
+            return model_stage
+        except RestException:
+            log.error("Model named %s with version %s does not exist", model_name, version)
+            return None
diff --git a/quadra/utils/patch/dataset.py b/quadra/utils/patch/dataset.py
index 82f03b39..3453b645 100644
--- a/quadra/utils/patch/dataset.py
+++ b/quadra/utils/patch/dataset.py
@@ -371,7 +371,6 @@ def __save_patch_dataset(
                             discard_in_training = False
 
                 if save_mask:
-
                     mask_name = f"{image_name}_{row_index * image_patches.shape[1] + col_index}{mask_extension}.png"
 
                     if discard_in_training:
@@ -1268,7 +1267,8 @@ def load_train_file(
     exclude_filter: Optional[List[str]] = None,
     class_to_skip: Optional[list] = None,
 ) -> Tuple[List[str], List[str]]:
-    """Load a train file and return a list of samples and a list of targets.
+    """Load a train file and return a list of samples and a list of targets. It is expected that train files will be in
+        the same location as the train_file_path.
 
     Args:
         train_file_path: Training file location
@@ -1305,6 +1305,9 @@ def load_train_file(
     samples = [samples[i] for i in valid_samples_indices]
     targets = [targets[i] for i in valid_samples_indices]
 
+    train_folder = os.path.dirname(train_file_path)
+    samples = [os.path.join(train_folder, x) for x in samples]
+
     return samples, targets
 
 
diff --git a/quadra/utils/utils.py b/quadra/utils/utils.py
index 0849d409..5fda5928 100644
--- a/quadra/utils/utils.py
+++ b/quadra/utils/utils.py
@@ -1,16 +1,18 @@
 """Common utility functions.
 Some of them are mostly based on https://github.com/ashleve/lightning-hydra-template.
 """
+import glob
 import json
 import logging
 import os
 import subprocess
 import sys
 import warnings
-from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, cast
 
 import cv2
 import dotenv
+import mlflow
 import numpy as np
 import pytorch_lightning as pl
 import rich.syntax
@@ -23,7 +25,11 @@
 from pytorch_lightning.utilities import rank_zero_only
 
 import quadra
-from quadra.callbacks.mlflow import check_file_server_dependencies, check_minio_credentials, get_mlflow_logger
+import quadra.utils.export as quadra_export
+from quadra.callbacks.mlflow import get_mlflow_logger
+from quadra.utils.mlflow import infer_signature_torch_model
+
+IMAGE_EXTENSIONS: List[str] = [".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".pbm", ".pgm", ".ppm", ".pxm", ".pnm"]
 
 
 def get_logger(name=__name__) -> logging.Logger:
@@ -73,13 +79,6 @@ def extras(config: DictConfig) -> None:
         if config.datamodule.get("num_workers"):
             config.datamodule.num_workers = 0
 
-    if config.core.get("upload_artifacts"):
-        if config.logger.get("mlflow") is not None and config.logger.mlflow.get("_target_") is not None:
-            if "http" in config.logger.mlflow.get("tracking_uri"):
-                log.info("Checking mlfow settings for artifact upload since <config.core.upload_artifacts=True>")
-                check_file_server_dependencies()
-                check_minio_credentials()
-
 
 @rank_zero_only
 def print_config(
@@ -222,44 +221,78 @@ def upload_file_tensorboard(file_path: str, tensorboard_logger: TensorBoardLogge
 
 def finish(
     config: DictConfig,
-    model: pl.LightningModule,
+    module: pl.LightningModule,
     datamodule: pl.LightningDataModule,
     trainer: pl.Trainer,
     callbacks: List[pl.Callback],
     logger: List[pl.loggers.Logger],
+    export_folder: str,
 ) -> None:
     """Upload config files to MLFlow server.
 
     Args:
         config: Configuration composed by Hydra.
-        model: LightningModule.
+        module: LightningModule.
         datamodule: LightningDataModule.
         trainer: LightningTrainer.
         callbacks: List of LightningCallbacks.
         logger: List of LightningLoggers.
+        export_folder: Folder where the deployment models are exported.
     """
     # pylint: disable=unused-argument
+
     if len(logger) > 0 and config.core.get("upload_artifacts"):
-        mflow_logger = get_mlflow_logger(trainer=trainer)
+        mlflow_logger = get_mlflow_logger(trainer=trainer)
         tensorboard_logger = get_tensorboard_logger(trainer=trainer)
+        file_names = ["config.yaml", "config_resolved.yaml", "config_tree.txt", "data/dataset.csv"]
 
-        if mflow_logger is not None:
+        if mlflow_logger is not None:
             config_paths = []
-            file_names = ["config.yaml", "config_resolved.yaml", "config_tree.txt"]
+
             for f in file_names:
                 if os.path.isfile(os.path.join(os.getcwd(), f)):
                     config_paths.append(os.path.join(os.getcwd(), f))
+
             for path in config_paths:
-                mflow_logger.experiment.log_artifact(
-                    run_id=mflow_logger.run_id, local_path=path, artifact_path="metadata"
+                mlflow_logger.experiment.log_artifact(
+                    run_id=mlflow_logger.run_id, local_path=path, artifact_path="metadata"
                 )
 
+            deployed_models = glob.glob(os.path.join(export_folder, "*"))
+            model_json: Optional[Dict[str, Any]] = None
+
+            if os.path.exists(os.path.join(export_folder, "model.json")):
+                with open(os.path.join(export_folder, "model.json"), "r") as json_file:
+                    model_json = json.load(json_file)
+
+            if model_json is not None:
+                for model_path in deployed_models:
+                    if model_path.endswith(".pt"):
+                        model, _ = quadra_export.import_deployment_model(model_path, device="cpu")
+
+                        input_size = model_json["input_size"]
+
+                        # Not a huge fan of this check
+                        if not isinstance(input_size[0], list):
+                            # Input size is not a list of lists
+                            input_size = [input_size]
+
+                        inputs = cast(List[Any], quadra_export.generate_torch_inputs(input_size, device="cpu"))
+                        signature = infer_signature_torch_model(model, inputs)
+
+                        with mlflow.start_run(run_id=mlflow_logger.run_id) as _:
+                            mlflow.pytorch.log_model(
+                                model,
+                                artifact_path=model_path,
+                                signature=signature,
+                            )
+
         if tensorboard_logger is not None:
             config_paths = []
-            file_names = ["config.yaml", "config_resolved.yaml", "config_tree.txt"]
             for f in file_names:
                 if os.path.isfile(os.path.join(os.getcwd(), f)):
                     config_paths.append(os.path.join(os.getcwd(), f))
+
             for path in config_paths:
                 upload_file_tensorboard(file_path=path, tensorboard_logger=tensorboard_logger)
 
diff --git a/quadra/utils/vit_explainability.py b/quadra/utils/vit_explainability.py
index 1247d307..bfaeb0b7 100644
--- a/quadra/utils/vit_explainability.py
+++ b/quadra/utils/vit_explainability.py
@@ -299,10 +299,9 @@ class LinearModelPytorchWrapper(torch.nn.Module):
 
     Args:
         backbone: Backbone
-        num_classes: Number of classes
         linear_classifier: ScikitLearn linear classifier model
-        device: The device to use. Defaults to "cpu"
         example_input: Input example needed to obtain output shape
+        device: The device to use. Defaults to "cpu"
     """
 
     def __init__(


Test patch:
diff --git a/quadra/utils/tests/models.py b/quadra/utils/tests/models.py
new file mode 100644
index 00000000..bc62d494
--- /dev/null
+++ b/quadra/utils/tests/models.py
@@ -0,0 +1,27 @@
+from typing import Any
+
+import torch
+from torch import nn
+
+
+class SingleInputModel(nn.Module):
+    """Model taking a single input."""
+
+    def forward(self, x: Any):
+        return x
+
+
+class DoubleInputModel(nn.Module):
+    """Model taking two inputs."""
+
+    def forward(self, x: Any, y: Any):
+        return x, y
+
+
+class UnsupportedInputModel(nn.Module):
+    """Model taking an unsupported input."""
+
+    def forward(self, x: torch.Tensor, y: str):
+        y = f"unsupported input: {y}"
+
+        return x
diff --git a/tests/datamodules/test_classification.py b/tests/datamodules/test_classification.py
index b188d7b3..00cb6292 100644
--- a/tests/datamodules/test_classification.py
+++ b/tests/datamodules/test_classification.py
@@ -139,15 +139,22 @@ def test_classification_patch_datamodule(classification_patch_dataset: classific
         info = PatchDatasetInfo(**json.load(f))
 
     # train samples are named like imagename_class.h5
-    train_samples_df = datamodule.data[datamodule.data["split"] == "train"]["samples"].tolist()
-    datamodule_train_samples = set(["_".join(s.split("_")[0:-1]) for s in train_samples_df])
+    train_samples_df = datamodule.train_data["samples"].tolist()
+    datamodule_train_samples = set(
+        [os.path.splitext(os.path.basename("_".join(s.split("_")[0:-1])))[0] for s in train_samples_df]
+    )
     # val samples are named like imagename_patchnumber.xyz
-    val_samples_df = datamodule.data[datamodule.data["split"] == "val"]["samples"].tolist()
-    datamodule_val_samples = set(["_".join(os.path.basename(s).split("_")[0:-1]) for s in val_samples_df])
+    val_samples_df = datamodule.val_data["samples"].tolist()
+    datamodule_val_samples = set(
+        [os.path.splitext("_".join(os.path.basename(s).split("_")[0:-1]))[0] for s in val_samples_df]
+    )
     # test samples are named like imagename_patchnumber.xyz and may contain #DISCARD# in the name
-    test_samples_df = datamodule.data[datamodule.data["split"] == "test"]["samples"].tolist()
+    test_samples_df = datamodule.test_data["samples"].tolist()
     datamodule_test_samples = set(
-        ["_".join(os.path.basename(s).replace("#DISCARD#", "").split("_")[0:-1]) for s in test_samples_df]
+        [
+            os.path.splitext("_".join(os.path.basename(s).replace("#DISCARD#", "").split("_")[0:-1]))[0]
+            for s in test_samples_df
+        ]
     )
 
     train_filenames = set([os.path.splitext(os.path.basename(s.image_path))[0] for s in info.train_files])
diff --git a/tests/models/test_signature_wrapper.py b/tests/models/test_signature_wrapper.py
new file mode 100644
index 00000000..70ea015b
--- /dev/null
+++ b/tests/models/test_signature_wrapper.py
@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+import torch
+
+from quadra.models.base import ModelSignatureWrapper
+from quadra.utils.export import generate_torch_inputs
+from quadra.utils.tests.models import DoubleInputModel, SingleInputModel, UnsupportedInputModel
+
+
+@torch.inference_mode()
+def test_simple_model():
+    """Test the input shape retrieval for a simple model."""
+    model = ModelSignatureWrapper(SingleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    model(x)
+    assert model.input_shapes == [(3, 224, 224)]
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+    model(*inputs)
+
+
+@torch.inference_mode()
+def test_double_input_model():
+    """Test the input shape retrieval for a model with multiple inputs."""
+    model = ModelSignatureWrapper(DoubleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    model(x, y)
+    assert model.input_shapes == [(3, 224, 224), (3, 448, 448)]
+
+    model.input_shapes = None
+    model(y=y, x=x)
+    # Check that the order of the inputs does not matter
+    assert model.input_shapes == [(3, 224, 224), (3, 448, 448)]
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+    assert len(inputs) == 2
+    assert inputs[0].shape == (1, 3, 224, 224)
+    assert inputs[1].shape == (1, 3, 448, 448)
+    model(*inputs)
+
+
+@torch.inference_mode()
+def test_dict_input_model():
+    """Test the input shape retrieval for a model with a dict input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    model({"x": x})
+    assert model.input_shapes == [{"x": (3, 224, 224)}]
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+    assert len(inputs) == 1
+    assert isinstance(inputs[0], dict)
+    assert inputs[0]["x"].shape == (1, 3, 224, 224)
+    model(*inputs)
+
+
+@torch.inference_mode()
+def test_tuple_input_model():
+    """Test the input shape retrieval for a model with a tuple input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    model((x, y))
+    assert model.input_shapes == [((3, 224, 224), (3, 448, 448))]
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+    assert len(inputs) == 1
+    assert len(inputs[0]) == 2
+    assert isinstance(inputs[0], tuple)
+    assert inputs[0][0].shape == (1, 3, 224, 224)
+    assert inputs[0][1].shape == (1, 3, 448, 448)
+    model(*inputs)
+
+
+@torch.inference_mode()
+def test_list_of_tensors_input_model():
+    """Test the input shape retrieval for a model with a list of tensors input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    model([x, y])
+    assert model.input_shapes == [[(3, 224, 224), (3, 448, 448)]]
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+    assert len(inputs) == 1
+    assert isinstance(inputs[0], list)
+    assert len(inputs[0]) == 2
+    assert inputs[0][0].shape == (1, 3, 224, 224)
+    assert inputs[0][1].shape == (1, 3, 448, 448)
+    model(*inputs)
+
+
+@torch.inference_mode()
+def test_list_of_dicts_input_model():
+    """Test the input shape retrieval for a model with a list of dicts input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    model([{"x": x}, {"y": y}])
+    assert model.input_shapes == [[{"x": (3, 224, 224)}, {"y": (3, 448, 448)}]]
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+    assert len(inputs) == 1
+    assert isinstance(inputs[0], list)
+    assert len(inputs[0]) == 2
+    assert inputs[0][0]["x"].shape == (1, 3, 224, 224)
+    assert inputs[0][1]["y"].shape == (1, 3, 448, 448)
+    model(*inputs)
+
+
+@torch.inference_mode()
+def test_tuple_of_dicts_input_model():
+    """Test the input shape retrieval for a model with a tuple of dicts input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    model(({"x": x}, {"y": y}))
+    assert model.input_shapes == [({"x": (3, 224, 224)}, {"y": (3, 448, 448)})]
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+    assert len(inputs) == 1
+    assert isinstance(inputs[0], tuple)
+    assert len(inputs[0]) == 2
+    assert inputs[0][0]["x"].shape == (1, 3, 224, 224)
+    assert inputs[0][1]["y"].shape == (1, 3, 448, 448)
+    model(*inputs)
+
+
+@torch.inference_mode()
+def test_unsupported_input_model():
+    """Test the input shape retrieval for a model with an unsupported input."""
+    model = ModelSignatureWrapper(UnsupportedInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = "test"
+    model(x, y)
+    assert model.input_shapes is None
+    assert model.disable
diff --git a/tests/tasks/test_anomaly.py b/tests/tasks/test_anomaly.py
index 4a4d327b..27a44549 100644
--- a/tests/tasks/test_anomaly.py
+++ b/tests/tasks/test_anomaly.py
@@ -77,7 +77,6 @@ def test_padim_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_datas
         f"datamodule.data_path={data_path}",
         "model.model.backbone=resnet18",
         f"model.dataset.task={task}",
-        "task.export_type=[torchscript]",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
@@ -105,7 +104,6 @@ def test_patchcore_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_d
         f"datamodule.data_path={data_path}",
         "model.model.backbone=resnet18",
         f"model.dataset.task={task}",
-        "task.export_type=[torchscript]",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
@@ -131,7 +129,7 @@ def test_cflow_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_datas
         f"datamodule.data_path={data_path}",
         "model.model.backbone=resnet18",
         f"model.dataset.task={task}",
-        "task.export_type=null",
+        "task.export_config=null",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
@@ -157,7 +155,6 @@ def test_csflow_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_data
         "experiment=base/anomaly/csflow",
         f"datamodule.data_path={data_path}",
         f"model.dataset.task={task}",
-        "task.export_type=[torchscript]",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
@@ -185,7 +182,6 @@ def test_fastflow_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_da
         f"datamodule.data_path={data_path}",
         "model.model.backbone=resnet18",
         f"model.dataset.task={task}",
-        "task.export_type=[torchscript]",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
@@ -212,7 +208,6 @@ def test_draem_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_datas
         "experiment=base/anomaly/draem",
         f"datamodule.data_path={data_path}",
         f"model.dataset.task={task}",
-        "task.export_type=[torchscript]",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
diff --git a/tests/utilities/test_mlflow_export.py b/tests/utilities/test_mlflow_export.py
new file mode 100644
index 00000000..4a875c42
--- /dev/null
+++ b/tests/utilities/test_mlflow_export.py
@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import importlib
+
+import pytest
+
+try:
+    importlib.import_module("mlflow")
+except ImportError:
+    pytest.skip("Mlflow is not installed", allow_module_level=True)
+
+from typing import Sequence
+
+import torch
+from mlflow.models import infer_signature
+from mlflow.types.schema import TensorSpec
+from pytest import raises
+
+from quadra.models.base import ModelSignatureWrapper
+from quadra.utils.export import generate_torch_inputs
+from quadra.utils.mlflow import infer_signature_input_torch, infer_signature_torch_model
+from quadra.utils.tests.models import DoubleInputModel, SingleInputModel
+
+
+def check_signature_equality(
+    input_signature: Sequence[TensorSpec] | dict[str, TensorSpec] | TensorSpec,
+    expected_signature: Sequence[TensorSpec] | dict[str, TensorSpec] | TensorSpec,
+):
+    """Assert that the generated signature matches the expected signature."""
+    is_equal = True
+
+    if isinstance(input_signature, Sequence) and isinstance(expected_signature, Sequence):
+        for c_input_signature, c_expected_signature in zip(input_signature, expected_signature):
+            is_equal = is_equal and check_signature_equality(c_input_signature, c_expected_signature)
+    elif isinstance(input_signature, dict) and isinstance(expected_signature, dict):
+        for k, v in input_signature.items():
+            is_equal = is_equal and check_signature_equality(v, expected_signature[k])
+    elif isinstance(input_signature, TensorSpec) and isinstance(expected_signature, TensorSpec):
+        is_equal = (
+            input_signature.type == expected_signature.type
+            and input_signature.shape == expected_signature.shape
+            and input_signature.name == expected_signature.name
+        )
+    else:
+        is_equal = False
+
+    return is_equal
+
+
+@torch.inference_mode()
+def test_single_tensor_signature():
+    """Test the input shape retrieval for a simple model."""
+    model = ModelSignatureWrapper(SingleInputModel())
+
+    x = torch.zeros(1, 3, 224, 224)
+    _ = model(x)
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+
+    signature = infer_signature_torch_model(model, inputs)
+
+    expected_input_signature = [TensorSpec(shape=(-1, *x.shape[1:]), type=x.numpy().dtype)]
+
+    # The model output is the same as the input
+    assert check_signature_equality(signature.inputs.inputs, expected_input_signature)
+    assert check_signature_equality(signature.outputs.inputs, expected_input_signature)
+
+
+@torch.inference_mode()
+def test_multiple_tensor_signature():
+    """Test the input shape retrieval for a model with multiple inputs."""
+    model = ModelSignatureWrapper(DoubleInputModel())
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    _ = model(x, y)
+    assert model.input_shapes == [(3, 224, 224), (3, 448, 448)]
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+
+    signature = infer_signature_torch_model(model, inputs)
+
+    expected_input_signature = [
+        TensorSpec(shape=(-1, *x.shape[1:]), type=x.numpy().dtype, name="output_0"),
+        TensorSpec(shape=(-1, *y.shape[1:]), type=y.numpy().dtype, name="output_1"),
+    ]
+
+    # The model output is the same as the input
+    assert check_signature_equality(signature.inputs.inputs, expected_input_signature)
+    assert check_signature_equality(signature.outputs.inputs, expected_input_signature)
+
+
+@torch.inference_mode()
+def test_dict_signature():
+    """Test the input shape retrieval for a model with a dict input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+
+    x = torch.zeros(1, 3, 224, 224)
+    _ = model({"x": x})
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+
+    signature = infer_signature_torch_model(model, inputs)
+
+    expected_input_signature = [TensorSpec(shape=(-1, *x.shape[1:]), type=x.numpy().dtype, name="x")]
+
+    # The model output is the same as the input
+    assert check_signature_equality(signature.inputs.inputs, expected_input_signature)
+    assert check_signature_equality(signature.outputs.inputs, expected_input_signature)
+
+
+@torch.inference_mode()
+def test_nested_tuple_signature():
+    """Test the input shape retrieval for a model with a tuple input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    outputs = model((x, y))
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+
+    expected_output_signature = [
+        TensorSpec(shape=(-1, *x.shape[1:]), type=x.numpy().dtype, name="output_0"),
+        TensorSpec(shape=(-1, *y.shape[1:]), type=y.numpy().dtype, name="output_1"),
+    ]
+
+    # Nested structures are not supported
+    with raises(ValueError):
+        infer_signature_input_torch(inputs)
+
+    signature = infer_signature(infer_signature_input_torch(outputs), infer_signature_input_torch(outputs))
+
+    assert check_signature_equality(signature.outputs.inputs, expected_output_signature)
+
+
+@torch.inference_mode()
+def test_nested_list_signature():
+    """Test the input shape retrieval for a model with a list of tensors input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    outputs = model([x, y])
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+
+    expected_output_signature = [
+        TensorSpec(shape=(-1, *x.shape[1:]), type=x.numpy().dtype, name="output_0"),
+        TensorSpec(shape=(-1, *y.shape[1:]), type=y.numpy().dtype, name="output_1"),
+    ]
+
+    # Nested structures are not supported
+    with raises(ValueError):
+        infer_signature_input_torch(inputs)
+
+    signature = infer_signature(infer_signature_input_torch(outputs), infer_signature_input_torch(outputs))
+
+    assert check_signature_equality(signature.outputs.inputs, expected_output_signature)
+
+
+@torch.inference_mode()
+def test_nested_dicts_signature():
+    """Test the input shape retrieval for a model with a list of dicts input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    outputs = model([{"x": x}, {"y": y}])
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+
+    # Nested structures are not supported
+    with raises(ValueError):
+        infer_signature_input_torch(inputs)
+
+    with raises(ValueError):
+        infer_signature_input_torch(outputs)
+
+
+@torch.inference_mode()
+def test_tuple_of_dicts_signature():
+    """Test the input shape retrieval for a model with a tuple of dicts input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    outputs = model(({"x": x}, {"y": y}))
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+
+    # Nested structures are not supported
+    with raises(ValueError):
+        infer_signature_input_torch(inputs)
+
+    with raises(ValueError):
+        infer_signature_input_torch(outputs)