Repository: orobix/quadra (language: python)

Issue:
Bugfix/git hyperparameter crash
## Summary

Right now we try to store the current commit inside the experiment hyperparameters for reproducibility, this causes an error when working inside of a git repository that doesn't contain any commit.

Resolves #26 

## Type of Change

- 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.)
- [X] I have added unit tests for my changes, or updated existing tests if necessary.
- [X] 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/.gitignore b/.gitignore
index 0c0710cb..695f945e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -68,3 +68,4 @@ docs/javascripts/images/*
 test-output.xml
 external/
 site/
+local/
diff --git a/CHANGELOG.md b/CHANGELOG.md
index cecde163..296dac33 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,47 @@
 # Changelog
 All notable changes to this project will be documented in this file.
 
+### [1.2.0]
+
+#### Added
+
+- Add plot_raw_outputs feature to class VisualizerCallback in anomaly detection, to save the raw images of the segmentation and heatmap output.
+- Add support for onnx exportation of trained models.
+- Add support for onnx model import in all evaluation tasks.
+- Add `export` configuration group to regulate exportation parameters.
+- Add `inference` configuration group to regulate inference parameters.
+- Add EfficientAD configuration for anomaly detection.
+- Add `acknowledgements` section to `README.md` file.
+- Add hashing parameters to datamodule configurations.
+
+#### Updated
+
+- Update anomalib library from version 0.4.0 to 0.7.0
+- Update mkdocs library from version 1.4.3 to 1.5.2
+- Update mkdocs-material library from version 9.1.18 to 9.2.8
+- Update mkdocstrings library by fixing the version to 0.23.0
+- Update mkdocs-material-extensions library by fixing the version to 1.1.1
+- Update mkdocs-autorefs library by fixing the version to 0.5.0
+- Update mkdocs-section-index library from version 0.3.5 to 0.3.6
+- Update mkdocstrings-python library from version 1.2.0 to 1.6.2
+- Update datamodule documentation for hashing.
+
+#### Changed 
+
+- Move `export_types` parameter from `task` configuration group to `export` configuration group under `types` parameter.
+- Refactor export model function to be more generic and be availble from the base task class.
+- Remove `save_backbone` parameter for scikit-learn based tasks.
+
+#### Fixed
+
+- Fix failures when trying to override `hydra` configuration groups due to wrong override order.
+- Fix certain anomalib models not loaded on the correct device.
+- Fix quadra crash when launching an experiment inside a git repository not fully initialized (e.g. without a single commit).
+- Fix documentation build failing due to wrong `mkdocstring` version.
+- Fix SSL docstrings 
+- Fix reference page URL to segmentation page in module management tutorial.
+- Fix `Makefile` command.
+
 ### [1.1.4]
 
 #### Fixed
diff --git a/Makefile b/Makefile
index 337ae43c..d5dff5fe 100644
--- a/Makefile
+++ b/Makefile
@@ -1,13 +1,16 @@
 # Makefile
 SHELL := /bin/bash
+DEVICE ?= cpu
 
 .PHONY: help
 help:
 	@echo "Commands:"
-	@echo "clean   		 : cleans all unnecessary files."
-	@echo "docs-serve    : serves the documentation."
-	@echo "docs-build    : builds the documentation."
-	@echo "style   		 : runs pre-commit."
+	@echo "clean				: cleans all unnecessary files."
+	@echo "docs-serve			: serves the documentation."
+	@echo "docs-build			: builds the documentation."
+	@echo "style				: runs pre-commit."
+	@echo "unit-tests:			: runs unit tests."
+	@echo "integration-tests:	: runs integration tests."
 
 # Cleaning
 .PHONY: clean
@@ -27,8 +30,13 @@ style:
 	pre-commit run --all --verbose
 .PHONY: docs-build
 docs-build:
-	 mkdocs build -d ./site
+	mkdocs build -d ./site
 
 .PHONY: docs-serve
 docs-serve:
-	 mkdocs serve
+	mkdocs serve
+
+.PHONY: unit-tests
+unit-tests:
+	@python -m pytest -v --disable-pytest-warnings --strict-markers --color=yes --device $(DEVICE)
+	
diff --git a/README.md b/README.md
index ba944fb9..1f9df34d 100644
--- a/README.md
+++ b/README.md
@@ -246,6 +246,21 @@ We rely on a combination of `Black`, `Pylint`, `Mypy`, `Ruff` and `Isort` to enf
 3. To run the webserver for real-time rendering and editing run `mkdocs serve` and visit `http://localhost:8000/`.
 4. If you want to export the static website to a specific folder  `mkdocs build -d <Destination Folder>`
 
+
+## Acknowledgements
+
+This project is based on many open-source libraries and frameworks, we would like to thank all the contributors for their work. Here is a list of the main libraries and frameworks we use:
+
+- [Pytorch](https://pytorch.org/) and [Pytorch Lightning](https://lightning.ai/) for training and deploying deep learning models. These two libraries are core part of training and testing tasks that allow us to run experiments on different devices in agile way.
+- Pretrained models are usually loaded from [Pytorch Hub](https://pytorch.org/hub/) or [Pytorch-image-models](https://github.com/huggingface/pytorch-image-models) (or called as `timm`).
+- Each specific task may rely on different libraries. For example, `segmentation` task uses [Segmentation_models.pytorch](https://github.com/qubvel/segmentation_models.pytorch) for loading backbones. The `anomaly detection` task uses a fork of [Anomalib](https://github.com/openvinotoolkit/anomalib) maintained by Orobix on [this repository](https://github.com/orobix/anomalib). We use light-weight ML models from [scikit-learn](https://scikit-learn.org/). We have also implementation of some SOTA models inside our library. 
+- Data processing and augmentation are done using [Albumentations](https://albumentations.ai/docs/) and [OpenCV](https://opencv.org/).
+- [Hydra](https://hydra.cc/docs/intro/) for composing configurations and running experiments. Hydra is a powerful framework that allows us to compose configurations from command line interface and run multiple experiments with different settings and hyperparameters. We have followed suggestions from `Configuring Experiments` section of [Hydra documentation](https://hydra.cc/docs/patterns/configuring_experiments/) and [lightning-hydra-template](https://github.com/ashleve/lightning-hydra-template) repository.
+- Documentation website is using [Material for MkDocs](https://squidfunk.github.io/mkdocs-material/) and [MkDocs](https://www.mkdocs.org/). For code documentation we are using [Mkdocstrings](https://mkdocstrings.github.io/). For releasing software versions we combine [Bumpver](https://github.com/mbarkhau/bumpver) and [Mike](https://github.com/jimporter/mike).
+- Models can be exported in different ways (`torchscript` or `torch` file). We have also added [ONNX](https://onnx.ai/) support for some models.
+- Testing framework is based on [Pytest](https://docs.pytest.org/en/) and related plug-ins.
+- Code quality is ensured by [pre-commit](https://pre-commit.com/) hooks. We are using [Black](https://github.com/psf/black) for formatting, [Pylint](https://www.pylint.org/) for linting, [Mypy](https://mypy.readthedocs.io/en/stable/) for type checking, [Isort](https://pycqa.github.io/isort/) for sorting imports, and [Ruff](https://github.com/astral-sh/ruff) for checking futher code and documentation quality.
+
 ## FAQ
 
 **How can I fix errors related to `GL` when I install full `opencv` package?**
diff --git a/docs/tutorials/configurations.md b/docs/tutorials/configurations.md
index bd442ff3..d0be0745 100644
--- a/docs/tutorials/configurations.md
+++ b/docs/tutorials/configurations.md
@@ -160,6 +160,9 @@ defaults:
   - override /scheduler: rop
   - override /transforms: default_resize
 
+export:
+  types: [torchscript]
+
 datamodule:
   num_workers: 8
   batch_size: 32
@@ -178,10 +181,6 @@ task:
   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/datamodules.md b/docs/tutorials/datamodules.md
index 5695ebfb..d5609560 100644
--- a/docs/tutorials/datamodules.md
+++ b/docs/tutorials/datamodules.md
@@ -155,3 +155,18 @@ To extend the base datamodule is necessary to implement the `_prepare_data` func
   - `split`: split of the image (train, test or val)
 
 These are generally the required fields, different tasks may require additional fields. For example, in the case of segmentation tasks, the `masks` field is required.
+
+
+
+## Data hashing
+
+During `prepare_data` call of each datamodule we apply hashing algorithm for each sample of the dataset. This information helps developer to track not only the data path used for the experiment but also to track the data content. This is useful when the data is stored in a remote location and the developer wants to check if the data is the same as the one used for the experiment. [BaseDataModule][quadra.datamodules.base.BaseDataModule] class has following arguments to control the hashing process:
+
+  - `enable_hashing`: If `True` the data will be hashed.
+  - `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 (`hash_type=size`) which is faster but less safe. Defaults to `content`.
+
+After the training is completed. The hash value of each sample used from given dataset will be saved under `hash` column inside `<experiment_folder>/data/dataset.csv` file.
+
+!!! info 
+  If the user wants to disable hashing from command line, it is possible to pass `datamodule.enable_hashing=False` as override argument.
\ No newline at end of file
diff --git a/docs/tutorials/devices_setup.md b/docs/tutorials/devices_setup.md
index 54556e65..3a8eea5e 100644
--- a/docs/tutorials/devices_setup.md
+++ b/docs/tutorials/devices_setup.md
@@ -38,13 +38,9 @@ _target_: quadra.tasks.SklearnClassification
 device: "cuda:0"
 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/anomaly_detection.md b/docs/tutorials/examples/anomaly_detection.md
index d3788f76..dde48111 100644
--- a/docs/tutorials/examples/anomaly_detection.md
+++ b/docs/tutorials/examples/anomaly_detection.md
@@ -90,6 +90,8 @@ experiments have shown that the previous models are better and faster to train.
 - [Fastflow](https://github.com/openvinotoolkit/anomalib/tree/main/src/anomalib/models/fastflow): FastFlow is a two-dimensional normalizing flow-based probability distribution estimator. It can be used as a plug-in module with any deep feature extractor, such as ResNet and vision transformer, for unsupervised anomaly detection and localisation. In the training phase, FastFlow learns to transform the input visual feature into a tractable distribution, and in the inference phase, it assesses the likelihood of identifying anomalies.
 - [DRAEM](https://github.com/openvinotoolkit/anomalib/tree/main/src/anomalib/models/draem): Is a reconstruction based algorithm that consists of a reconstructive subnetwork and a discriminative subnetwork. DRAEM is trained on simulated anomaly images, generated by augmenting normal input images from the training set with a random Perlin noise mask extracted from an unrelated source of image data. The reconstructive subnetwork is an autoencoder architecture that is trained to reconstruct the original input images from the augmented images. The reconstructive submodel is trained using a combination of L2 loss and Structural Similarity loss. The input of the discriminative subnetwork consists of the channel-wise concatenation of the (augmented) input image and the output of the reconstructive subnetwork. The output of the discriminative subnetwork is an anomaly map that contains the predicted anomaly scores for each pixel location. The discriminative subnetwork is trained using Focal Loss
 - [CS-FLOW](https://github.com/openvinotoolkit/anomalib/tree/main/src/anomalib/models/csflow): The central idea of the paper is to handle fine-grained representations by incorporating global and local image context. This is done by taking multiple scales when extracting features and using a fully-convolutional normalizing flow to process the scales jointly.
+- [EfficientAd](https://github.com/openvinotoolkit/anomalib/tree/main/src/anomalib/models/efficient_ad)
+Fast anomaly segmentation algorithm that consists of a distilled pre-trained teacher model, a student model and an autoencoder. It detects local anomalies via the teacher-student discrepany and global anomalies via the student-autoencoder discrepancy.
 
 For a detailed description of the models and their parameters please refer to the anomalib documentation.
 
@@ -122,6 +124,8 @@ callbacks:
     output_path: anomaly_output
     threshold_type: ${callbacks.min_max_normalization.threshold_type}
     disable: true
+    plot_only_wrong: false
+    plot_raw_outputs: false
 ```
 
 The min_max_normalization callback is used to normalize the anomaly maps to the range [0, 1] such that the threshold will become 0.5. 
@@ -132,6 +136,11 @@ The post processing configuration allow to specify the method used to compute th
 
 The visualizer callback is used to produce a visualization of the results on the test data, when the min_max_normalization callback is used the input_are_normalized flag must be set to true and the threshold_type should match the one used for normalization. By default it is disabled as it may take a while to compute, to enable just set `disable: false`.
 
+In the context where many images are supplied to our model, we may be more interested in restricting the output images that are generated to only the cases where the result is not correct. By default it is disabled, to enable just set `plot_only_wrong: true`.
+
+The display of the outputs of a model can be done in a preset format. However, this option may not be as desired, or may be affecting the resolution of the images. In order to give more flexibility to the generation of reports, the heatmap and segmentation ouput files can be generated independently and with the same resolution of the original image. By default it is disabled, to enable just set `plot_raw_outputs: true`.
+
+
 ### Anomalib configuration
 Anomalib library doesn't use hydra but still uses yaml configurations that are found under `model/anomalib`.
 This for example is the configuration used for PADIM.
@@ -177,7 +186,7 @@ As already mentioned anomaly detection requires just good images for training, t
 
 ### Experiment
 
-Suppose that we want to run the experiment on the given dataset using the PADIM technique. We can define take the generic padim config for mnist as an example found under `experiment/generic/mnist/anomaly/padim.yaml`.
+Suppose that we want to run the experiment on the given dataset using the PADIM technique. We can take the generic padim config for mnist as an example found under `experiment/generic/mnist/anomaly/padim.yaml`.
 
 ```yaml
 # @package _global_
@@ -185,6 +194,9 @@ defaults:
   - base/anomaly/padim
   - override /datamodule: generic/mnist/anomaly/base
 
+export:
+  types: [torchscript]
+
 model:
   model:
     input_size: [224, 224]
@@ -219,7 +231,7 @@ trainer:
   check_val_every_n_epoch: ${trainer.max_epochs}
 ```
 
-We start from the base configuration for PADIM, then we override the datamodule to use the generic mnist datamodule. Using this configuration we specify that we want to use PADIM, extracting features using the resnet18 backbone with image size 224x224, the dataset is `mnist`, we specify that the task is taken from the anomalib configuration which specify it to be segmentation. One very important thing to watch out is the `check_val_every_n_epoch` parameter. This parameter should match the number of epochs for `PADIM` and `Patchcore`, the reason is that in the validation phase the model will be fitted and we want the fit to be done only once and on all the data, increasing the max_epoch is useful when we apply data augmentation, otherwise it doesn't make a lot of sense as we would fit the model on the same, replicated data.
+We start from the base configuration for PADIM, then we override the datamodule to use the generic mnist datamodule. Using this configuration we specify that we want to use PADIM, extracting features using the resnet18 backbone with image size 224x224, the dataset is `mnist`, we specify that the task is taken from the anomalib configuration which specify it to be segmentation. One very important thing to watch out is the `check_val_every_n_epoch` parameter. This parameter should match the number of epochs for `PADIM` and `Patchcore`, the reason is that in the validation phase the model will be fitted and we want the fit to be done only once and on all the data, increasing the max_epoch is useful when we apply data augmentation, otherwise it doesn't make a lot of sense as we would fit the model on the same, replicated data. The model will be exported at the end of the training phase, as we have specified the `export.types` parameter to `torchscript` the model will be exported only in torchscript format.
 
 ### Run
 
@@ -282,7 +294,7 @@ task:
 
 By default, the inference will recompute the threshold based on test data to maximize the F1-score, if you want to use the threshold from the training phase you can set the `use_training_threshold` parameter to true.
 
-The model path is the path to an exported model, at the moment only `torchscript` models are supported (exported automatically after a training experiment). Right now only the `CFLOW` model is not supported for inference as it's not compatible with torchscript.
+The model path is the path to an exported model, at the moment `torchscript` and `onnx` models are supported (exported automatically after a training experiment). Right now only the `CFLOW` model is not supported for inference as it's not compatible with botyh torchscript and onnx.
 
 An inference configuration using the mnist dataset is found under `configs/experiment/generic/mnist/anomaly/inference.yaml`.
 
diff --git a/docs/tutorials/examples/classification.md b/docs/tutorials/examples/classification.md
index 4c21705a..2cdff54b 100644
--- a/docs/tutorials/examples/classification.md
+++ b/docs/tutorials/examples/classification.md
@@ -114,9 +114,6 @@ task:
   report: True
   output:
     example: True
-  export_config:
-    types: [pytorch, torchscript]
-    input_shapes: # Redefine the input shape if not automatically inferred
 
 
 core:
@@ -172,6 +169,9 @@ defaults:
   - override /backbone: vit16_tiny
   - _self_
 
+export:
+  types: [onnx, torchscript]
+
 datamodule:
   num_workers: 12
   batch_size: 32
@@ -187,9 +187,6 @@ task:
   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:
@@ -236,7 +233,7 @@ checkpoints           config_tree.txt  deployment_model  test
 config_resolved.yaml  data             main.log
 ```
 
-Where `checkpoints` contains the pytorch lightning checkpoints of the model, `data` contains the joblib dump of the datamodule with its parameters and dataset split, `deployment_model` contains the model in exported format (default is torchscript), `test` contains the test artifacts.
+Where `checkpoints` contains the pytorch lightning checkpoints of the model, `data` contains the joblib dump of the datamodule with its parameters and dataset split, `deployment_model` contains the model in exported format (in this case onnx and torchscript, but by default is only torchscript), `test` contains the test artifacts.
 
 ## Evaluation
 
diff --git a/docs/tutorials/examples/multilabel_classification.md b/docs/tutorials/examples/multilabel_classification.md
index cd9fe0b6..883afcbf 100644
--- a/docs/tutorials/examples/multilabel_classification.md
+++ b/docs/tutorials/examples/multilabel_classification.md
@@ -139,9 +139,6 @@ task:
   report: False 
   output:
     example: False 
-  export_config:
-    types: [torchscript]
-    input_shapes: # Redefine the input shape if not automatically inferred
 
   
 logger:
diff --git a/docs/tutorials/examples/segmentation.md b/docs/tutorials/examples/segmentation.md
index 5aa8fb43..a153e08b 100644
--- a/docs/tutorials/examples/segmentation.md
+++ b/docs/tutorials/examples/segmentation.md
@@ -162,6 +162,9 @@ defaults:
   - base/segmentation/smp_multiclass  # use smp file as default
   - _self_ # use this file as final config
 
+export:
+  types: [onnx, torchscript]
+
 backbone:
   model:
     classes: 4 # The total number of classes (background + foreground)
@@ -171,10 +174,6 @@ task:
   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
@@ -200,7 +199,7 @@ core:
     When defining the `idx_to_class` dictionary, the keys should be the class index and the values should be the class name. The class index starts from 1.
 
 
-In the final configuration experiment we have specified the path to the dataset, batch size, split files, GPU device, experiment name and toggled some evaluation options.
+In the final configuration experiment we have specified the path to the dataset, batch size, split files, GPU device, experiment name and toggled some evaluation options, moreover we have specified that we want to export the model to `onnx` and `torchscript` formats.
 
 By default data will be logged to `Mlflow`. If `Mlflow` is not available it's possible to configure a simple csv logger by adding an override to the file above:
   
diff --git a/docs/tutorials/examples/sklearn_classification.md b/docs/tutorials/examples/sklearn_classification.md
index 7c14bc26..588ec191 100644
--- a/docs/tutorials/examples/sklearn_classification.md
+++ b/docs/tutorials/examples/sklearn_classification.md
@@ -95,17 +95,14 @@ defaults:
   - override /trainer: sklearn_classification
   - override /datamodule: base/sklearn_classification
 
+export:
+  types: [pytorch, torchscript]
+  
 backbone:
   model:
     pretrained: true
     freeze: true
 
-task:
-  export_config:
-    types: [pytorch, torchscript]
-    input_shapes: # Redefine the input shape if not automatically inferred
-
-
 core:
   tag: "run"
   name: "sklearn-classification"
@@ -121,6 +118,7 @@ By default the experiment will use dino_vitb8 as backbone, resizing the images t
 It will also export the model in two formats, "torchscript" and "pytorch".
 
 An actual configuration file based on the above could be this one (suppose it's saved under `configs/experiment/custom_experiment/sklearn_classification.yaml`):
+
 ```yaml
 # @package _global_
 
@@ -132,6 +130,9 @@ defaults:
 core:
   name: experiment-name
 
+export:
+  types: [pytorch, torchscript]
+
 datamodule:
   data_path: path_to_dataset
   batch_size: 64
@@ -148,18 +149,13 @@ task:
   device: cuda:0
   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_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).
+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.types` parameter can be used to export the model in different formats, at the moment `torchscript`, `onnx` and `pytorch` are supported.
+The backbone (in torchscript and pytorch 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
 
@@ -181,7 +177,7 @@ classification_experiment_2  config_tree.txt              test
 
 Each `classification_experiment_X` folder contains the metrics for the corresponding fold while the `classification_experiment` folder contains the metrics computed aggregating the results of all the folds.
 
-The `data` folder contains a joblib version of the datamodule containing parameters and splits for reproducibility. The `deployment_model` folder contains the backbone exported in torchscript format if `save_backbone` to true alongside the joblib version of trained classifier. The `test` folder contains the metrics for the final test on all the data after the model has been trained on both train and validation.
+The `data` folder contains a joblib version of the datamodule containing parameters and splits for reproducibility. The `deployment_model` folder contains the backbone exported in torchscript and pytorch format alongside the joblib version of trained classifier. The `test` folder contains the metrics for the final test on all the data after the model has been trained on both train and validation.
 
 ## Evaluation
 The same datamodule specified before can be used for inference by setting the `phase` parameter to `test`. 
diff --git a/docs/tutorials/examples/sklearn_patch_classification.md b/docs/tutorials/examples/sklearn_patch_classification.md
index 35db4453..d1580afa 100644
--- a/docs/tutorials/examples/sklearn_patch_classification.md
+++ b/docs/tutorials/examples/sklearn_patch_classification.md
@@ -205,6 +205,9 @@ defaults:
   - override /backbone: resnet18
   - _self_
 
+export:
+  types: [torchscript]
+  
 core:
   name: experiment-name
 
@@ -222,14 +225,9 @@ task:
   device: cuda:2
   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
-
 ```
 
 This will train a resnet18 model on the given dataset, using 256 as batch size and skipping the background class during training.
@@ -253,7 +251,7 @@ config_tree.txt                  main.log
 
 Inside the `classification_patch_experiment` folder you should find some report utilities computed over the validation dataset, like the confusion matrix. The `reconstruction_results.json` file contains the reconstruction metrics computed over the validation dataset in terms of covered defects, it will also contain the coordinates of the polygons extracted over predicted areas of the image with the same label.
 
-The `data` folder contains a joblib version of the datamodule containing parameters and splits for reproducibility. The `deployment_model` folder contains the backbone exported in torchscript format if `save_backbone` to true alongside the joblib version of trained classifier.
+The `data` folder contains a joblib version of the datamodule containing parameters and splits for reproducibility. The `deployment_model` folder contains the backbone exported in torchscript format alongside the joblib version of trained classifier.
 
 ## Evaluation
 The same datamodule specified before can be used for inference. 
diff --git a/docs/tutorials/examples/ssl.md b/docs/tutorials/examples/ssl.md
index 836796e5..13c41666 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_config.types` 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 `export.types` parameter (default `torchscript`). 
 
 ### Run (Advanced) - Changing transformations
 
diff --git a/docs/tutorials/export.md b/docs/tutorials/export.md
new file mode 100644
index 00000000..61a563f6
--- /dev/null
+++ b/docs/tutorials/export.md
@@ -0,0 +1,95 @@
+# Export models for inference
+
+In this section we will see how quadra allows you to export your trained models for inference. We will see how to export models for both Lightning based tasks and Sklearn based tasks.
+By default the standard export format is Torchscript, but you can also export models to ONNX or plain Pytorch (this is done for particular operations like gradcam).
+
+## Standard export configuration
+
+The standard configuration for exporting models is located at `configs/export/default.yaml` and it is shown below:
+
+```yaml
+types: [torchscript]
+input_shapes: # Redefine the input shape if not automatically inferred
+onnx:
+  # torch.onnx.export options
+  input_names: # If null automatically inferred
+  output_names: # If null automatically inferred
+  dynamic_axes: # If null automatically inferred
+  export_params: true
+  opset_version: 16
+  do_constant_folding: true
+  # Custom options
+  fixed_batch_size: # If not null export with fixed batch size (ignore dynamic axes)
+  simplify: true
+```
+
+`types` is a list of the export types that you want to perform. The available types are `torchscript`, `onnx` and `pytorch`. By default the models will be saved under the `deployment_model` folder of the experiment with extension `.pt`, `.onnx` and `.pth` respectively. Pytorch models will be saved alongside the yaml configuration for the model itself so that you can easily load them back in python.
+
+`input_shapes` is a parameter that will be `None` most of the time, quadra features a model wrapper that is capable of inferring the input shape of the trained model based on its forward function. It supports a large variety of custom forward functions where parameters are combinations of lists, tuples or dicts. However, if the model wrapper is not able to infer the input shape you can specify it here, the format is a list of tuples/lists/dicts where each element represents a single input shape without batch size. For example if your model has an input of shape (1, 3, 224, 224) you can specify it as:
+
+```yaml
+input_shapes:
+  - [3, 224, 224]
+```
+
+Onnx support a set of extra options passed as kwargs to the `torch.onnx.export`, once again we will try to automatically infer most of them but if you need to specify them you can do it under the `onnx` section of the configuration file. The two custom options are:
+
+- `fixed_batch_size`: It can be used to export the model with a fixed batch size, this is useful if you want to use the model in a context where you know the batch size in advance. If you specify this option the model will be exported with a fixed batch size and the dynamic axes will be ignored.
+- `simplify`: If true the model will be simplified using the [onnx-simplifier](https://github.com/daquexian/onnx-simplifier) package, the resulting model is called `model_simplified.onnx` and it is saved alongside the original model.
+
+## Lightning based tasks
+
+Currently quadra supports exporting models for the following tasks:
+
+- Image classification
+- Image segmentation
+- Anomaly detection (certain models may not be supported)
+- SSL training
+
+## Sklearn based tasks
+
+Currently quadra supports exporting models for the following tasks:
+
+- Image classification
+- Image classification with patches
+
+When working with sklearn based tasks alongside the exported backbone in the `deployment_model` folder you will also find a `classifier.joblib` containing the exported sklearn model.
+
+## Importing models for quadra evaluation
+
+Quadra exported models are fully compatible with quadra evaluation tasks, this is possible because quadra uses a model wrapper emulating the standard pytorch interface for all the exported models. 
+
+### Standard inference configuration
+
+Evaluation models are regulated by a configuration file located at `configs/inference/default.yaml` shown below:
+
+```yaml
+onnx:
+  session_options:
+    inter_op_num_threads: 8
+    intra_op_num_threads: 8
+    graph_optimization_level:
+      _target_: onnxruntime.GraphOptimizationLevel
+      value: 99 # ORT_ENABLE_ALL
+    enable_mem_pattern: true
+    enable_cpu_mem_arena: true
+    enable_profiling: false
+    enable_mem_reuse: true
+    execution_mode:
+      _target_: onnxruntime.ExecutionMode
+      value: 0 # ORT_SEQUENTIAL
+    execution_order:
+      _target_: onnxruntime.ExecutionOrder
+      value: 0 # DEFAULT
+    log_severity_level: 2
+    log_verbosity_level: 0
+    logid: ""
+    optimized_model_filepath: ""
+    use_deterministic_compute: false
+    profile_file_prefix: onnxruntime_profile_
+
+pytorch:
+torchscript:
+```
+
+Right now we support custom option only for ONNX runtime, but we plan to add more inference configuration options in the future.
\ No newline at end of file
diff --git a/docs/tutorials/model_management.md b/docs/tutorials/model_management.md
index 50b514be..d7be160d 100644
--- a/docs/tutorials/model_management.md
+++ b/docs/tutorials/model_management.md
@@ -19,7 +19,7 @@ By defining an abstract class, we establish a common interface that can be imple
 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.
+    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:
 
diff --git a/mkdocs.yml b/mkdocs.yml
index ae89f3d2..4e661e0c 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -15,6 +15,7 @@ nav:
       - Setting up devices: tutorials/devices_setup.md
       - Reproducibility: tutorials/reproducibility.md
       - Model Management: tutorials/model_management.md
+      - Export models for inference: tutorials/export.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 e2c20ed3..0c0b95d8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,33 +7,33 @@ name = "quadra"
 version = "1.1.4"
 description = "Deep Learning experiment orchestration library"
 authors = [
-  {name = "Alessandro Polidori", email = "alessandro.polidori@orobix.com"},
-  {name = "Federico Belotti", email = "federico.belotti@orobix.com"},
-  {name = "Lorenzo Mammana", email = "lorenzo.mammana@orobix.com"},
-  {name = "Refik Can Malli", email = "refikcan.malli@orobix.com"},
-  {name = "Silvia Bianchetti", email = "silvia.bianchetti@orobix.com"},
+  { name = "Alessandro Polidori", email = "alessandro.polidori@orobix.com" },
+  { name = "Federico Belotti", email = "federico.belotti@orobix.com" },
+  { name = "Lorenzo Mammana", email = "lorenzo.mammana@orobix.com" },
+  { name = "Refik Can Malli", email = "refikcan.malli@orobix.com" },
+  { name = "Silvia Bianchetti", email = "silvia.bianchetti@orobix.com" },
 ]
 keywords = ["deep learning", "experiment", "lightning", "hydra-core"]
-license = {file = "LICENSE"}
-readme = {file = "README.md", content-type = "text/markdown"}
+license = { file = "LICENSE" }
+readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.8,<3.10"
 classifiers = [
-    "Programming Language :: Python :: 3",
-    "Intended Audience :: Developers",
-    "Intended Audience :: Education",
-    "Intended Audience :: Science/Research",
-    "Topic :: Scientific/Engineering :: Artificial Intelligence",
-    "Topic :: Software Development",
-    "Topic :: Software Development :: Libraries",
-    "Topic :: Software Development :: Libraries :: Python Modules",
-    "License :: OSI Approved :: Apache Software License",
+  "Programming Language :: Python :: 3",
+  "Intended Audience :: Developers",
+  "Intended Audience :: Education",
+  "Intended Audience :: Science/Research",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+  "Topic :: Software Development",
+  "Topic :: Software Development :: Libraries",
+  "Topic :: Software Development :: Libraries :: Python Modules",
+  "License :: OSI Approved :: Apache Software License",
 ]
 
 dependencies = [
   # --------- pytorch --------- #
   "torch==1.13.1",
   "torchvision==0.14.1",
-  "torchmetrics==0.10.*", # constrained by anomalib
+  "torchmetrics==0.10.*",           # constrained by anomalib
   "torchsummary==1.5.*",
   "pytorch-lightning>=1.9.1,<1.10",
   # --------- hydra --------- #
@@ -42,11 +42,11 @@ dependencies = [
   "hydra-optuna-sweeper==1.2.*",
   # --------- loggers --------- #
   "mlflow==2.3.1",
-  "boto3==1.26.*", # needed for artifact storage
-  "minio==7.1.*", # needed for artifact storage
+  "boto3==1.26.*",       # needed for artifact storage
+  "minio==7.1.*",        # needed for artifact storage
   "tensorboard==2.11.*",
   # --------- others --------- #
-  "Pillow==9.3.0", # required by label-studio-converter
+  "Pillow==9.3.0",                                                        # required by label-studio-converter
   "pandas==1.1.*",
   "opencv-python-headless==4.7.0.*",
   "python-dotenv==0.21.*",
@@ -61,46 +61,47 @@ dependencies = [
   "scikit-multilearn==0.2.*",
   "tripy==1.0.*",
   "h5py==3.8.*",
-  "timm==0.6.12", # required by smp
+  "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.7.0+obx.1.2.0",
   "xxhash==3.2.*",
 ]
 
 [project.optional-dependencies]
-test = [
-    "pytest==7.2.*",
-    "pytest-cov==4.0.*",
-]
+test = ["pytest==7.2.*", "pytest-cov==4.0.*", "pytest-lazy-fixture==0.6.*"]
 
 dev = [
-    "interrogate==1.5.*",
-    "black==22.12.*",
-    "isort==5.11.*",
-    "pre-commit==3.0.*",
-    "pylint==2.16.*",
-    "bump2version==1.0.*",
-    "types-PyYAML==6.0.12.*",
-    "mypy==1.0.*",
-    "ruff==0.0.257",
-    "pandas-stubs==1.5.3.*",
-    "bumpver==2023.1124",
-    "twine==4.0.*",
-    "build==0.10.*"
-
+  "interrogate==1.5.*",
+  "black==22.12.*",
+  "isort==5.11.*",
+  "pre-commit==3.0.*",
+  "pylint==2.16.*",
+  "bump2version==1.0.*",
+  "types-PyYAML==6.0.12.*",
+  "mypy==1.0.*",
+  "ruff==0.0.257",
+  "pandas-stubs==1.5.3.*",
+  "bumpver==2023.1124",
+  "twine==4.0.*",
+  "build==0.10.*",
 ]
 
 docs = [
-    "mkdocs==1.4.3",
-    "mkdocs-material==9.1.18",
-    "mkdocstrings-python==1.2.0",
-    "mkdocs-gen-files==0.5.0",
-    "mkdocs-literate-nav==0.6.0",
-    "mkdocs-section-index==0.3.5",
-    "mike==1.1.2",
-    "cairosvg==2.7.0"
+  "mkdocs==1.5.2",
+  "mkdocs-literate-nav==0.6.0",
+  "mkdocs-section-index==0.3.6",
+  "mkdocstrings==0.23.0",
+  "mkdocs-autorefs==0.5.0",
+  "mkdocs-gen-files==0.5.0",
+  "mkdocs-material==9.2.8",
+  "mkdocstrings-python==1.6.2",
+  "mkdocs-material-extensions==1.1.1",
+  "mike==1.1.2",
+  "cairosvg==2.7.0",
 ]
 
+onnx = ["onnx==1.14.0", "onnxsim==0.4.28", "onnxruntime-gpu==1.15.0"]
+
 [tool.setuptools]
 include-package-data = true
 
@@ -119,10 +120,10 @@ repository = "https://github.com/orobix/quadra"
 [tool.bumpver]
 current_version = "1.1.4"
 version_pattern = "MAJOR.MINOR.PATCH"
-commit_message  = "build: Bump version {old_version} -> {new_version}"
-commit          = true
-tag             = false
-push            = false
+commit_message = "build: Bump version {old_version} -> {new_version}"
+commit = true
+tag = false
+push = false
 
 [tool.bumpver.file_patterns]
 "pyproject.toml" = ['current_version = "{version}"', 'version = "{version}"']
@@ -162,9 +163,7 @@ skip_gitignore = true
 testpaths = ["tests"]
 python_files = "test_*.py"
 addopts = "--strict-markers --disable-pytest-warnings"
-markers = [
-    "slow: marks tests as slow (deselect with '-m \"not slow\"')",
-]
+markers = ["slow: marks tests as slow (deselect with '-m \"not slow\"')"]
 
 # Pytest coverage
 [tool.coverage.run]
@@ -188,8 +187,8 @@ quiet = false
 whitelist_regex = []
 color = true
 ignore_regex = [
-  "^get$", 
-  "^mock_.*", 
+  "^get$",
+  "^mock_.*",
   ".*BaseClass.*",
   ".*on_train.*",
   ".*on_validation.*",
@@ -243,25 +242,20 @@ warn_return_any = false
 exclude = ["quadra/utils/tests", "tests"]
 
 [tool.ruff]
-select = [
-    "D",
-]
+select = ["D"]
 ignore = [
-    "D100",
-    # this is controlled by interrogate with exlude_regex
-    # we can skip it here
-    "D102", 
-    "D104",
-    "D105",
-    "D107",
-    # no blank line after summary line. This might be not required.
-    # usually we violate this rule
-    "D205"
-]
-exclude = [
-    "Makefile",
-    ".gitignore",
+  "D100",
+  # this is controlled by interrogate with exlude_regex
+  # we can skip it here
+  "D102",
+  "D104",
+  "D105",
+  "D107",
+  # no blank line after summary line. This might be not required.
+  # usually we violate this rule
+  "D205",
 ]
+exclude = ["Makefile", ".gitignore"]
 
 [tool.ruff.pydocstyle]
 convention = "google"
diff --git a/quadra/callbacks/anomalib.py b/quadra/callbacks/anomalib.py
index 3f5bbce7..c9ab213a 100644
--- a/quadra/callbacks/anomalib.py
+++ b/quadra/callbacks/anomalib.py
@@ -1,6 +1,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 
+import cv2
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
@@ -92,6 +93,7 @@ class VisualizerCallback(Callback):
         threshold_type: Either 'pixel' or 'image'. If 'pixel', the threshold is computed on the pixel-level.
         disable: whether to disable the callback.
         plot_only_wrong: whether to plot only the images that are not correctly predicted.
+        plot_raw_outputs: Saves the raw images of the segmentation and heatmap output.
     """
 
     def __init__(
@@ -102,6 +104,7 @@ def __init__(
         threshold_type: str = "pixel",
         disable: bool = False,
         plot_only_wrong: bool = False,
+        plot_raw_outputs: bool = False,
     ) -> None:
         self.inputs_are_normalized = inputs_are_normalized
         self.output_path = output_path
@@ -109,6 +112,7 @@ def __init__(
         self.disable = disable
         self.task = task
         self.plot_only_wrong = plot_only_wrong
+        self.plot_raw_outputs = plot_raw_outputs
 
     def _add_images(self, visualizer: Visualizer, filename: Path, output_label_folder: str):
         """Save image to logger/local storage.
@@ -214,9 +218,26 @@ def on_test_batch_end(
             visualizer.figure.suptitle(
                 f"F1 threshold: {threshold}, Mask_max: {anomaly_map.max():.3f}, Anomaly_score: {anomaly_score:.3f}"
             )
-            self._add_images(visualizer, Path(filename), output_label_folder)
+            filename = Path(filename)
+            self._add_images(visualizer, filename, output_label_folder)
             visualizer.close()
 
+            if self.plot_raw_outputs:
+                for raw_output, raw_name in zip([heat_map, vis_img], ["heatmap", "segmentation"]):
+                    if raw_name == "segmentation":
+                        raw_output = (raw_output * 255).astype(np.uint8)
+                    raw_output = cv2.cvtColor(raw_output, cv2.COLOR_RGB2BGR)
+                    raw_filename = (
+                        Path(self.output_path)
+                        / "images"
+                        / output_label_folder
+                        / filename.parent.name
+                        / "raw_outputs"
+                        / Path(filename.stem + f"_{raw_name}.png")
+                    )
+                    raw_filename.parent.mkdir(parents=True, exist_ok=True)
+                    cv2.imwrite(str(raw_filename), raw_output)
+
     def on_test_end(self, _trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
         """Sync logs.
 
diff --git a/quadra/configs/callbacks/default_anomalib.yaml b/quadra/configs/callbacks/default_anomalib.yaml
index 6627eacd..f7fa301f 100644
--- a/quadra/configs/callbacks/default_anomalib.yaml
+++ b/quadra/configs/callbacks/default_anomalib.yaml
@@ -18,6 +18,8 @@ visualizer:
   output_path: anomaly_output
   threshold_type: ${callbacks.min_max_normalization.threshold_type}
   disable: true
+  plot_only_wrong: false
+  plot_raw_outputs: false
 # Standard callbacks
 early_stopping:
   _target_: pytorch_lightning.callbacks.EarlyStopping
diff --git a/quadra/configs/config.yaml b/quadra/configs/config.yaml
index 9cf9f795..97a47a6b 100644
--- a/quadra/configs/config.yaml
+++ b/quadra/configs/config.yaml
@@ -15,9 +15,12 @@ defaults:
   - datamodule: null
   - callbacks: default
   - logger: mlflow
-  - experiment: null
+  - export: default
+  - inference: default
   - hparams_search: null
   - hydra: default
+  # BE CAREFUL TO KEEP EXPERIMENT LAST
+  - experiment: null
 
   # enable color logging
   - override hydra/hydra_logging: colorlog
diff --git a/quadra/configs/datamodule/base/anomaly.yaml b/quadra/configs/datamodule/base/anomaly.yaml
index 4536f914..c7473d86 100644
--- a/quadra/configs/datamodule/base/anomaly.yaml
+++ b/quadra/configs/datamodule/base/anomaly.yaml
@@ -11,3 +11,6 @@ val_transform: ${transforms.val_transform}
 phase: train
 valid_area_mask:
 crop_area:
+enable_hashing: true
+hash_size: 64
+hash_type: content
diff --git a/quadra/configs/datamodule/base/classification.yaml b/quadra/configs/datamodule/base/classification.yaml
index 1f3e8443..9c8b65cd 100644
--- a/quadra/configs/datamodule/base/classification.yaml
+++ b/quadra/configs/datamodule/base/classification.yaml
@@ -16,3 +16,6 @@ name:
 dataset:
   _target_: hydra.utils.get_method
   path: quadra.datasets.classification.ClassificationDataset
+enable_hashing: true
+hash_size: 64
+hash_type: content
diff --git a/quadra/configs/datamodule/base/multilabel_classification.yaml b/quadra/configs/datamodule/base/multilabel_classification.yaml
index 028e4b1c..84373cbc 100644
--- a/quadra/configs/datamodule/base/multilabel_classification.yaml
+++ b/quadra/configs/datamodule/base/multilabel_classification.yaml
@@ -18,3 +18,6 @@ train_transform: ${transforms.train_transform}
 test_transform: ${transforms.test_transform}
 val_transform: ${transforms.val_transform}
 class_to_idx: null
+enable_hashing: true
+hash_size: 64
+hash_type: content
diff --git a/quadra/configs/datamodule/base/segmentation.yaml b/quadra/configs/datamodule/base/segmentation.yaml
index 191b20f0..35c72cb6 100644
--- a/quadra/configs/datamodule/base/segmentation.yaml
+++ b/quadra/configs/datamodule/base/segmentation.yaml
@@ -13,3 +13,6 @@ test_split_file:
 val_split_file:
 num_data_class:
 exclude_good: false
+enable_hashing: true
+hash_size: 64
+hash_type: content
diff --git a/quadra/configs/datamodule/base/segmentation_multiclass.yaml b/quadra/configs/datamodule/base/segmentation_multiclass.yaml
index 68cc677d..2aefa695 100644
--- a/quadra/configs/datamodule/base/segmentation_multiclass.yaml
+++ b/quadra/configs/datamodule/base/segmentation_multiclass.yaml
@@ -15,3 +15,6 @@ val_split_file:
 exclude_good: false
 num_data_train:
 one_hot_encoding:
+enable_hashing: true
+hash_size: 64
+hash_type: content
diff --git a/quadra/configs/datamodule/base/sklearn_classification.yaml b/quadra/configs/datamodule/base/sklearn_classification.yaml
index 7452ef69..44cc68ee 100644
--- a/quadra/configs/datamodule/base/sklearn_classification.yaml
+++ b/quadra/configs/datamodule/base/sklearn_classification.yaml
@@ -18,3 +18,6 @@ cache: false
 limit_training_data:
 train_split_file:
 test_split_file:
+enable_hashing: true
+hash_size: 64
+hash_type: content
diff --git a/quadra/configs/datamodule/base/sklearn_classification_patch.yaml b/quadra/configs/datamodule/base/sklearn_classification_patch.yaml
index 2aeca838..269ac0ad 100644
--- a/quadra/configs/datamodule/base/sklearn_classification_patch.yaml
+++ b/quadra/configs/datamodule/base/sklearn_classification_patch.yaml
@@ -12,3 +12,6 @@ test_transform: ${transforms.test_transform}
 val_transform: ${transforms.val_transform}
 balance_classes: false
 class_to_skip_training:
+enable_hashing: true
+hash_size: 64
+hash_type: content
diff --git a/quadra/configs/datamodule/base/ssl.yaml b/quadra/configs/datamodule/base/ssl.yaml
index 67285a85..9beaddef 100644
--- a/quadra/configs/datamodule/base/ssl.yaml
+++ b/quadra/configs/datamodule/base/ssl.yaml
@@ -16,3 +16,6 @@ val_size: 0.3
 test_size: 0.1
 split_validation: true
 class_to_idx:
+enable_hashing: true
+hash_size: 64
+hash_type: content
diff --git a/quadra/configs/experiment/base/anomaly/efficient_ad.yaml b/quadra/configs/experiment/base/anomaly/efficient_ad.yaml
new file mode 100644
index 00000000..4d1fd57a
--- /dev/null
+++ b/quadra/configs/experiment/base/anomaly/efficient_ad.yaml
@@ -0,0 +1,76 @@
+# @package _global_
+
+defaults:
+  - override /datamodule: base/anomaly
+  - override /model: anomalib/efficient_ad
+  - override /optimizer: null
+  - override /scheduler: null
+  - override /transforms: default_resize
+  - override /loss: null
+  - override /task: anomalib/efficient_ad
+  - override /backbone: null
+  - override /trainer: lightning_gpu
+  - override /callbacks: default_anomalib
+  - _self_
+
+datamodule:
+  num_workers: 12
+  train_batch_size: 32
+  test_batch_size: 32
+  task: ${model.dataset.task}
+  category:
+  phase: train
+
+core:
+  tag: "run"
+  name: efficient_ad_${model.model.model_size}_${trainer.max_epochs}
+  test_after_training: true
+
+logger:
+  mlflow:
+    experiment_name:
+    run_name: ${core.name}
+
+trainer:
+  devices: [2]
+  accelerator: auto
+  strategy:
+  accumulate_grad_batches: 1
+  amp_backend: native
+  auto_lr_find: false
+  auto_scale_batch_size: false
+  auto_select_gpus: false
+  benchmark: false
+  check_val_every_n_epoch: ${trainer.max_epochs}
+  default_root_dir: null
+  detect_anomaly: false
+  deterministic: false
+  enable_checkpointing: true
+  enable_model_summary: true
+  enable_progress_bar: true
+  fast_dev_run: false
+  gradient_clip_val: 0
+  ipus: null
+  limit_predict_batches: 1.0
+  limit_test_batches: 1.0
+  limit_train_batches: 1.0
+  limit_val_batches: 1.0
+  log_every_n_steps: 50
+  max_epochs: 20
+  max_steps: 20000
+  max_time: null
+  min_epochs: null
+  min_steps: null
+  move_metrics_to_cpu: false
+  multiple_trainloader_mode: max_size_cycle
+  num_nodes: 1
+  num_sanity_val_steps: 0
+  overfit_batches: 0.0
+  plugins: null
+  precision: 32
+  profiler: null
+  replace_sampler_ddp: true
+  sync_batchnorm: false
+  tpu_cores: null
+  track_grad_norm: -1
+  val_check_interval: 1.0 # Don't validate before extracting features.
diff --git a/quadra/configs/experiment/base/anomaly/padim.yaml b/quadra/configs/experiment/base/anomaly/padim.yaml
index f9b1964b..1b05f7a2 100644
--- a/quadra/configs/experiment/base/anomaly/padim.yaml
+++ b/quadra/configs/experiment/base/anomaly/padim.yaml
@@ -64,7 +64,6 @@ trainer:
   move_metrics_to_cpu: false
   multiple_trainloader_mode: max_size_cycle
   num_nodes: 1
-  num_processes: 1
   num_sanity_val_steps: 0
   overfit_batches: 0.0
   plugins: null
diff --git a/quadra/configs/experiment/base/anomaly/patchcore.yaml b/quadra/configs/experiment/base/anomaly/patchcore.yaml
index 1430f35c..f71ec1e2 100644
--- a/quadra/configs/experiment/base/anomaly/patchcore.yaml
+++ b/quadra/configs/experiment/base/anomaly/patchcore.yaml
@@ -64,7 +64,6 @@ trainer:
   move_metrics_to_cpu: false
   multiple_trainloader_mode: max_size_cycle
   num_nodes: 1
-  num_processes: 1
   num_sanity_val_steps: 0
   overfit_batches: 0.0
   plugins: null
diff --git a/quadra/configs/experiment/base/classification/classification.yaml b/quadra/configs/experiment/base/classification/classification.yaml
index 3ce38c9b..99616028 100644
--- a/quadra/configs/experiment/base/classification/classification.yaml
+++ b/quadra/configs/experiment/base/classification/classification.yaml
@@ -9,6 +9,9 @@ defaults:
   - override /scheduler: rop
   - override /transforms: default_resize
 
+export:
+  types: [torchscript, pytorch]
+
 datamodule:
   num_workers: 8
   batch_size: 32
@@ -24,9 +27,6 @@ model:
 task:
   lr_multiplier: 0.0
   run_test: True
-  export_config:
-    types: [torchscript, pytorch]
-    input_shapes:
   report: True
   gradcam: True
   output:
diff --git a/quadra/configs/experiment/base/classification/multilabel_classification.yaml b/quadra/configs/experiment/base/classification/multilabel_classification.yaml
index 2dcb5bca..8ae16870 100644
--- a/quadra/configs/experiment/base/classification/multilabel_classification.yaml
+++ b/quadra/configs/experiment/base/classification/multilabel_classification.yaml
@@ -9,6 +9,9 @@ defaults:
   - override /scheduler: rop
   - override /transforms: default_resize
 
+export:
+  types: [torchscript, pytorch]
+
 datamodule:
   num_workers: 8
   batch_size: 32
diff --git a/quadra/configs/experiment/base/classification/sklearn_classification.yaml b/quadra/configs/experiment/base/classification/sklearn_classification.yaml
index 1a833b99..fca8785c 100644
--- a/quadra/configs/experiment/base/classification/sklearn_classification.yaml
+++ b/quadra/configs/experiment/base/classification/sklearn_classification.yaml
@@ -8,16 +8,14 @@ defaults:
   - override /trainer: sklearn_classification
   - override /datamodule: base/sklearn_classification
 
+export:
+  types: [torchscript, pytorch]
+
 backbone:
   model:
     pretrained: true
     freeze: true
 
-task:
-  export_config:
-    types: [torchscript, pytorch]
-    input_shapes: null
-
 core:
   tag: "run"
   name: "sklearn-classification"
diff --git a/quadra/configs/experiment/default.yaml b/quadra/configs/experiment/default.yaml
index 33bb8ecc..44148d86 100644
--- a/quadra/configs/experiment/default.yaml
+++ b/quadra/configs/experiment/default.yaml
@@ -12,4 +12,4 @@ defaults:
   - override /trainer: null
   - override /callbacks: null
   - override /logger: null
-  - _self_
\ No newline at end of file
+  - _self_
diff --git a/quadra/configs/experiment/generic/mvtec/anomaly/efficient_ad.yaml b/quadra/configs/experiment/generic/mvtec/anomaly/efficient_ad.yaml
new file mode 100644
index 00000000..0abf6601
--- /dev/null
+++ b/quadra/configs/experiment/generic/mvtec/anomaly/efficient_ad.yaml
@@ -0,0 +1,38 @@
+# @package _global_
+defaults:
+  - base/anomaly/efficient_ad
+  - override /datamodule: generic/mvtec/anomaly/base
+
+transforms:
+  input_height: 256
+  input_width: 256
+
+datamodule:
+  num_workers: 12
+  train_batch_size: 1
+  test_batch_size: 1
+  task: ${model.dataset.task}
+  category: hazelnut
+
+callbacks:
+  min_max_normalization:
+    threshold_type: "pixel"
+
+print_config: false
+
+core:
+  tag: "run"
+  test_after_training: true
+  upload_artifacts: true
+  name: efficient_ad_${datamodule.category}_${model.model.model_size}_${trainer.max_epochs}
+
+logger:
+  mlflow:
+    experiment_name: mvtec-anomaly
+    run_name: ${core.name}
+
+trainer:
+  devices: [0]
+  max_epochs: 50
+  max_steps: 20000
+  check_val_every_n_epoch: ${trainer.max_epochs}
diff --git a/quadra/configs/export/default.yaml b/quadra/configs/export/default.yaml
new file mode 100644
index 00000000..02ca1e1f
--- /dev/null
+++ b/quadra/configs/export/default.yaml
@@ -0,0 +1,13 @@
+types: [torchscript]
+input_shapes: # Redefine the input shape if not automatically inferred
+onnx:
+  # torch.onnx.export options
+  input_names: # If null automatically inferred
+  output_names: # If null automatically inferred
+  dynamic_axes: # If null automatically inferred
+  export_params: true
+  opset_version: 16
+  do_constant_folding: true
+  # Custom options
+  fixed_batch_size: # If not null export with fixed batch size (ignore dynamic axes)
+  simplify: true
diff --git a/quadra/configs/inference/default.yaml b/quadra/configs/inference/default.yaml
new file mode 100644
index 00000000..a0746d8e
--- /dev/null
+++ b/quadra/configs/inference/default.yaml
@@ -0,0 +1,26 @@
+onnx:
+  session_options:
+    inter_op_num_threads: 8
+    intra_op_num_threads: 8
+    graph_optimization_level:
+      _target_: onnxruntime.GraphOptimizationLevel
+      value: 99 # ORT_ENABLE_ALL
+    enable_mem_pattern: true
+    enable_cpu_mem_arena: true
+    enable_profiling: false
+    enable_mem_reuse: true
+    execution_mode:
+      _target_: onnxruntime.ExecutionMode
+      value: 0 # ORT_SEQUENTIAL
+    execution_order:
+      _target_: onnxruntime.ExecutionOrder
+      value: 0 # DEFAULT
+    log_severity_level: 2
+    log_verbosity_level: 0
+    logid: ""
+    optimized_model_filepath: ""
+    use_deterministic_compute: false
+    profile_file_prefix: onnxruntime_profile_
+
+pytorch:
+torchscript:
diff --git a/quadra/configs/model/anomalib/cfa.yaml b/quadra/configs/model/anomalib/cfa.yaml
index 0e959d67..c564370c 100644
--- a/quadra/configs/model/anomalib/cfa.yaml
+++ b/quadra/configs/model/anomalib/cfa.yaml
@@ -32,7 +32,7 @@ trainer:
   gradient_clip_val: 0
   gradient_clip_algorithm: norm
   num_nodes: 1
-  devices: 1
+  devices: [0]
   enable_progress_bar: true
   overfit_batches: 0.0
   track_grad_norm: -1
diff --git a/quadra/configs/model/anomalib/csflow.yaml b/quadra/configs/model/anomalib/csflow.yaml
index 66899ebe..85c78f82 100644
--- a/quadra/configs/model/anomalib/csflow.yaml
+++ b/quadra/configs/model/anomalib/csflow.yaml
@@ -66,7 +66,6 @@ trainer:
   move_metrics_to_cpu: false
   multiple_trainloader_mode: max_size_cycle
   num_nodes: 1
-  num_processes: null
   num_sanity_val_steps: 0
   overfit_batches: 0.0
   plugins: null
diff --git a/quadra/configs/model/anomalib/efficient_ad.yaml b/quadra/configs/model/anomalib/efficient_ad.yaml
new file mode 100644
index 00000000..1f270b7d
--- /dev/null
+++ b/quadra/configs/model/anomalib/efficient_ad.yaml
@@ -0,0 +1,30 @@
+dataset:
+  task: segmentation
+
+model:
+  name: efficientad
+  teacher_out_channels: 384
+  model_size: small # options: [small, medium]
+  lr: 0.0001
+  image_size: [256, 256]
+  weight_decay: 0.00001
+  padding: false
+  pad_maps: true # relevant for "padding: false", see EfficientAd in lightning_model.py
+  # generic params
+  normalization_method: min_max # options: [null, min_max, cdf]
+  train_batch_size: 1 # ${datamodule.train_batch_size}
+  pretrained_models_dir: ${oc.env:HOME}/.quadra/models/efficient_ad
+  imagenette_dir: ${oc.env:HOME}/.quadra/datasets/
+  pretrained_teacher_type: nelson
+
+metrics:
+  image:
+    - F1Score
+    - AUROC
+  pixel:
+    - F1Score
+    - AUROC
+  threshold:
+    method: adaptive # options: [adaptive, manual]
+    manual_image: null
+    manual_pixel: null
diff --git a/quadra/configs/model/anomalib/padim.yaml b/quadra/configs/model/anomalib/padim.yaml
index 7b9a38cd..275daaf2 100644
--- a/quadra/configs/model/anomalib/padim.yaml
+++ b/quadra/configs/model/anomalib/padim.yaml
@@ -9,6 +9,7 @@ dataset:
     random_tile_count: 16
 
 model:
+  name: padim
   input_size: [224, 224]
   backbone: resnet18
   layers:
diff --git a/quadra/configs/task/anomalib/cfa.yaml b/quadra/configs/task/anomalib/cfa.yaml
index f675715c..9082c1ed 100644
--- a/quadra/configs/task/anomalib/cfa.yaml
+++ b/quadra/configs/task/anomalib/cfa.yaml
@@ -3,7 +3,3 @@ report: true
 run_test: true
 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 9cd0912b..bd5a34f3 100644
--- a/quadra/configs/task/anomalib/cflow.yaml
+++ b/quadra/configs/task/anomalib/cflow.yaml
@@ -3,7 +3,3 @@ report: true
 run_test: true
 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 68a1ec24..a94bced7 100644
--- a/quadra/configs/task/anomalib/csflow.yaml
+++ b/quadra/configs/task/anomalib/csflow.yaml
@@ -3,7 +3,3 @@ report: true
 run_test: true
 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 09d29e51..dcdface7 100644
--- a/quadra/configs/task/anomalib/draem.yaml
+++ b/quadra/configs/task/anomalib/draem.yaml
@@ -3,7 +3,3 @@ report: true
 run_test: true
 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/efficient_ad.yaml b/quadra/configs/task/anomalib/efficient_ad.yaml
new file mode 100644
index 00000000..31311228
--- /dev/null
+++ b/quadra/configs/task/anomalib/efficient_ad.yaml
@@ -0,0 +1,5 @@
+_target_: quadra.tasks.anomaly.AnomalibDetection
+report: true
+run_test: true
+module_function:
+  _target_: anomalib.models.efficient_ad.lightning_model.EfficientAdLightning
diff --git a/quadra/configs/task/anomalib/fastflow.yaml b/quadra/configs/task/anomalib/fastflow.yaml
index d681d71d..4885f584 100644
--- a/quadra/configs/task/anomalib/fastflow.yaml
+++ b/quadra/configs/task/anomalib/fastflow.yaml
@@ -3,7 +3,3 @@ report: true
 run_test: true
 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 25deb46d..ed79a11e 100644
--- a/quadra/configs/task/anomalib/padim.yaml
+++ b/quadra/configs/task/anomalib/padim.yaml
@@ -3,7 +3,3 @@ report: true
 run_test: true
 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 4dfdc376..417acbc2 100644
--- a/quadra/configs/task/anomalib/patchcore.yaml
+++ b/quadra/configs/task/anomalib/patchcore.yaml
@@ -3,7 +3,3 @@ report: true
 run_test: true
 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/anomaly.yaml b/quadra/configs/task/anomaly.yaml
deleted file mode 100644
index 3e78d264..00000000
--- a/quadra/configs/task/anomaly.yaml
+++ /dev/null
@@ -1,1 +0,0 @@
-_target_: quadra.tasks.AnomalyDetection
diff --git a/quadra/configs/task/classification.yaml b/quadra/configs/task/classification.yaml
index a2ecb8e3..c022f447 100644
--- a/quadra/configs/task/classification.yaml
+++ b/quadra/configs/task/classification.yaml
@@ -4,6 +4,3 @@ 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 51b0b2e3..13d9ee92 100644
--- a/quadra/configs/task/segmentation.yaml
+++ b/quadra/configs/task/segmentation.yaml
@@ -7,6 +7,3 @@ 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 feae7377..8e2ec122 100644
--- a/quadra/configs/task/sklearn_classification.yaml
+++ b/quadra/configs/task/sklearn_classification.yaml
@@ -2,10 +2,6 @@ _target_: quadra.tasks.SklearnClassification
 device: "cuda:0"
 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 2a25c366..ba3f73d4 100644
--- a/quadra/configs/task/sklearn_classification_patch.yaml
+++ b/quadra/configs/task/sklearn_classification_patch.yaml
@@ -2,10 +2,6 @@ _target_: quadra.tasks.PatchSklearnClassification
 device: cuda:2
 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 24ae88e4..fe2b5c9d 100644
--- a/quadra/configs/task/ssl.yaml
+++ b/quadra/configs/task/ssl.yaml
@@ -1,5 +1,2 @@
 _target_: quadra.tasks.SSL
 run_test: false
-export_config:
-  types: [torchscript]
-  input_shapes:
diff --git a/quadra/datamodules/classification.py b/quadra/datamodules/classification.py
index 6dfcebc4..f7aa3e94 100644
--- a/quadra/datamodules/classification.py
+++ b/quadra/datamodules/classification.py
@@ -47,6 +47,7 @@ class ClassificationDataModule(BaseDataModule):
         val_split_file: The file with validation split. Defaults to None.
         test_split_file: The file with test split. Defaults to None.
         class_to_idx: The mapping from class name to index. Defaults to None.
+        **kwargs: Additional arguments for BaseDataModule.
     """
 
     def __init__(
@@ -75,6 +76,7 @@ def __init__(
         test_split_file: Optional[str] = None,
         val_split_file: Optional[str] = None,
         class_to_idx: Optional[Dict[str, int]] = None,
+        **kwargs: Any,
     ):
         super().__init__(
             data_path=data_path,
@@ -90,6 +92,7 @@ def __init__(
             n_aug_to_take=n_aug_to_take,
             replace_str_from=replace_str_from,
             replace_str_to=replace_str_to,
+            **kwargs,
         )
         self.replace_str = None
         self.exclude_filter = exclude_filter
diff --git a/quadra/datasets/ssl.py b/quadra/datasets/ssl.py
index 5ee2b49d..14e74b5d 100644
--- a/quadra/datasets/ssl.py
+++ b/quadra/datasets/ssl.py
@@ -72,10 +72,11 @@ class TwoSetAugmentationDataset(Dataset):
         num_local_transforms: Number of local transformations to apply. In total you will have
             two + num_local_transforms transformations for each image. First element of the array will always
             return the original image.
-            images[0] = global_transform[0](original_image)
-            images[1] = global_transform[1](original_image)
-            images[2:] = local_transform(s)(original_image)
-            ...
+
+    Example:
+        >>> images[0] = global_transform[0](original_image)
+        >>> images[1] = global_transform[1](original_image)
+        >>> images[2:] = local_transform(s)(original_image)
     """
 
     def __init__(
diff --git a/quadra/models/base.py b/quadra/models/base.py
index 173e16e4..f9557f1c 100644
--- a/quadra/models/base.py
+++ b/quadra/models/base.py
@@ -36,7 +36,7 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
                 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! "
+                    "provide the input shapes manually setting the config.export.input_shapes parameter! "
                     "Alternatively you could try to use a forward with supported input types (and their compositions) "
                     "(list, tuple, dict, tensors)."
                 )
diff --git a/quadra/models/evaluation.py b/quadra/models/evaluation.py
new file mode 100644
index 00000000..f40c25a6
--- /dev/null
+++ b/quadra/models/evaluation.py
@@ -0,0 +1,268 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, cast
+
+import numpy as np
+import torch
+from hydra.utils import instantiate
+from omegaconf import DictConfig, OmegaConf
+from torch import nn
+from torch.jit import RecursiveScriptModule
+
+from quadra.utils.logger import get_logger
+
+try:
+    import onnxruntime as ort  # noqa
+
+    ONNX_AVAILABLE = True
+except ImportError:
+    ONNX_AVAILABLE = False
+
+
+log = get_logger(__name__)
+
+
+class BaseEvaluationModel(ABC):
+    """Base interface for all evaluation models."""
+
+    def __init__(self, config: DictConfig) -> None:
+        self.model: Any
+        self.model_path: str | None
+        self.device: str
+        self.config = config
+
+    @abstractmethod
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        pass
+
+    @abstractmethod
+    def load_from_disk(self, model_path: str, device: str = "cpu"):
+        """Load model from disk."""
+
+    @abstractmethod
+    def to(self, device: str):
+        """Move model to device."""
+
+    @abstractmethod
+    def eval(self):
+        """Set model to evaluation mode."""
+
+    @abstractmethod
+    def half(self):
+        """Convert model to half precision."""
+
+    @abstractmethod
+    def cpu(self):
+        """Move model to cpu."""
+
+    @property
+    def training(self) -> bool:
+        """Return whether model is in training mode."""
+        return False
+
+
+class TorchscriptEvaluationModel(BaseEvaluationModel):
+    """Wrapper for torchscript models."""
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self.model(*args, **kwargs)
+
+    def load_from_disk(self, model_path: str, device: str = "cpu"):
+        """Load model from disk."""
+        self.model_path = model_path
+        self.device = device
+
+        model = cast(RecursiveScriptModule, torch.jit.load(self.model_path))
+        model.eval()
+        model.to(self.device)
+
+        self.model = model
+
+    def to(self, device: str):
+        """Move model to device."""
+        self.model.to(device)
+        self.device = device
+
+    def eval(self):
+        """Set model to evaluation mode."""
+        self.model.eval()
+
+    @property
+    def training(self) -> bool:
+        """Return whether model is in training mode."""
+        return self.model.training
+
+    def half(self):
+        """Convert model to half precision."""
+        self.model.half()
+
+    def cpu(self):
+        """Move model to cpu."""
+        self.model.cpu()
+
+
+class TorchEvaluationModel(TorchscriptEvaluationModel):
+    """Wrapper for torch models.
+
+    Args:
+        model_architecture: Optional torch model architecture
+    """
+
+    def __init__(self, config: DictConfig, model_architecture: nn.Module) -> None:
+        super().__init__(config=config)
+        self.model = model_architecture
+        self.model.eval()
+        device = next(self.model.parameters()).device
+        self.device = str(device)
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self.model(*args, **kwargs)
+
+    def load_from_disk(self, model_path: str, device: str = "cpu"):
+        """Load model from disk."""
+        self.model_path = model_path
+        self.device = device
+        self.model.load_state_dict(torch.load(self.model_path))
+        self.model.eval()
+        self.model.to(self.device)
+
+
+class ONNXEvaluationModel(BaseEvaluationModel):
+    """Wrapper for ONNX models. It's designed to provide a similar interface to standard torch models."""
+
+    def __init__(self, config: DictConfig) -> None:
+        if not ONNX_AVAILABLE:
+            raise ImportError(
+                "onnxruntime is not installed. Please install ONNX capabilities for quadra with: pip install .[onnx]"
+            )
+        super().__init__(config=config)
+        self.session_options = self.generate_session_options()
+
+    def generate_session_options(self) -> ort.SessionOptions:
+        """Generate session options from the current config."""
+        session_options = ort.SessionOptions()
+
+        if hasattr(self.config, "session_options") and self.config.session_options is not None:
+            session_options_dict = cast(
+                dict[str, Any], OmegaConf.to_container(self.config.session_options, resolve=True)
+            )
+            for key, value in session_options_dict.items():
+                if isinstance(value, dict) and "_target_" in value:
+                    value = instantiate(value)
+
+                setattr(session_options, key, value)
+
+        return session_options
+
+    def __call__(self, *inputs: np.ndarray | torch.Tensor) -> Any:
+        """Run inference on the model and return the output as torch tensors."""
+        # TODO: Maybe we can support also kwargs
+        use_pytorch = False
+
+        onnx_inputs: dict[str, np.ndarray | torch.Tensor] = {}
+
+        for onnx_input, current_input in zip(self.model.get_inputs(), inputs):
+            if isinstance(current_input, torch.Tensor):
+                onnx_inputs[onnx_input.name] = current_input
+                use_pytorch = True
+            elif isinstance(current_input, np.ndarray):
+                onnx_inputs[onnx_input.name] = current_input
+            else:
+                raise ValueError(f"Invalid input type: {type(inputs)}")
+
+            if use_pytorch and isinstance(current_input, np.ndarray):
+                raise ValueError("Cannot mix torch and numpy inputs")
+
+        if use_pytorch:
+            onnx_output = self._forward_from_pytorch(cast(dict[str, torch.Tensor], onnx_inputs))
+        else:
+            onnx_output = self._forward_from_numpy(cast(dict[str, np.ndarray], onnx_inputs))
+
+        onnx_output = [torch.from_numpy(x).to(self.device) if isinstance(x, np.ndarray) else x for x in onnx_output]
+
+        if len(onnx_output) == 1:
+            onnx_output = onnx_output[0]
+
+        return onnx_output
+
+    def _forward_from_pytorch(self, input_dict: dict[str, torch.Tensor]):
+        """Run inference on the model and return the output as torch tensors."""
+        io_binding = self.model.io_binding()
+        device_type = self.device.split(":")[0]
+
+        for k, v in input_dict.items():
+            if not v.is_contiguous():
+                # If not contiguous onnx give wrong results
+                v = v.contiguous()
+
+            io_binding.bind_input(
+                name=k,
+                device_type=device_type,
+                # Weirdly enough onnx wants 0 for cpu
+                device_id=0 if device_type == "cpu" else int(self.device.split(":")[1]),
+                element_type=np.float32,
+                shape=tuple(v.shape),
+                buffer_ptr=v.data_ptr(),
+            )
+
+        for x in self.model.get_outputs():
+            # TODO: Is it possible to also bind the output? We require info about output dimensions
+            io_binding.bind_output(name=x.name)
+
+        self.model.run_with_iobinding(io_binding)
+
+        output = io_binding.copy_outputs_to_cpu()
+
+        return output
+
+    def _forward_from_numpy(self, input_dict: dict[str, np.ndarray]):
+        """Run inference on the model and return the output as numpy array."""
+        ort_outputs = [x.name for x in self.model.get_outputs()]
+
+        onnx_output = self.model.run(ort_outputs, input_dict)
+
+        return onnx_output
+
+    def load_from_disk(self, model_path: str, device: str = "cpu"):
+        """Load model from disk."""
+        self.model_path = model_path
+        self.device = device
+
+        ort_providers = self._get_providers(device)
+        self.model = ort.InferenceSession(self.model_path, providers=ort_providers, sess_options=self.session_options)
+
+    def _get_providers(self, device: str) -> list[tuple[str, dict[str, Any]] | str]:
+        """Return the providers for the ONNX model based on the device."""
+        ort_providers: list[tuple[str, dict[str, Any]] | str]
+
+        if device == "cpu":
+            ort_providers = ["CPUExecutionProvider"]
+        else:
+            ort_providers = [
+                (
+                    "CUDAExecutionProvider",
+                    {
+                        "device_id": int(device.split(":")[1]),
+                    },
+                )
+            ]
+
+        return ort_providers
+
+    def to(self, device: str):
+        """Move model to device."""
+        self.device = device
+        ort_providers = self._get_providers(device)
+        self.model.set_providers(ort_providers)
+
+    def eval(self):
+        """Fake interface to match torch models."""
+
+    def half(self):
+        """Convert model to half precision."""
+        raise NotImplementedError("At the moment ONNX models do not support half method.")
+
+    def cpu(self):
+        """Move model to cpu."""
+        self.to("cpu")
diff --git a/quadra/modules/classification/base.py b/quadra/modules/classification/base.py
index c68685bb..affc7bbf 100644
--- a/quadra/modules/classification/base.py
+++ b/quadra/modules/classification/base.py
@@ -135,10 +135,14 @@ def prepare_gradcam(self) -> None:
         """Instantiate gradcam handlers."""
         if isinstance(self.model.features_extractor, timm.models.resnet.ResNet):
             target_layers = [cast(BaseNetworkBuilder, self.model).features_extractor.layer4[-1]]  # type: ignore[index]
+
+            # Get model current device
+            device = next(self.model.parameters()).device
+
             self.cam = GradCAM(
                 model=self.model,
                 target_layers=target_layers,
-                use_cuda=torch.cuda.is_available(),
+                use_cuda=device.type == "cuda",
             )
             # Activating gradients
             for p in self.model.features_extractor.layer4[-1].parameters():
diff --git a/quadra/tasks/anomaly.py b/quadra/tasks/anomaly.py
index 3f36958a..38375b46 100644
--- a/quadra/tasks/anomaly.py
+++ b/quadra/tasks/anomaly.py
@@ -24,7 +24,7 @@
 from quadra.tasks.base import Evaluation, LightningTask
 from quadra.utils import utils
 from quadra.utils.classification import get_results
-from quadra.utils.export import export_torchscript_model
+from quadra.utils.export import export_model
 
 log = utils.get_logger(__name__)
 
@@ -41,11 +41,6 @@ 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_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__(
@@ -55,14 +50,12 @@ def __init__(
         checkpoint_path: Optional[str] = None,
         run_test: bool = True,
         report: bool = True,
-        export_config: Optional[DictConfig] = None,
     ):
         super().__init__(
             config=config,
             checkpoint_path=checkpoint_path,
             run_test=run_test,
             report=report,
-            export_config=export_config,
         )
         self._module: AnomalyModule
         self.module_function = module_function
@@ -112,43 +105,31 @@ def prepare(self) -> None:
 
     def export(self) -> None:
         """Export model for production."""
-        if self.export_config is None or len(self.export_config.types) == 0:
-            log.info("No export type specified skipping export")
-            return
-
         if self.config.trainer.get("fast_dev_run"):
             log.warning("Skipping export since fast_dev_run is enabled")
             return
 
         model = self.module.model
 
-        input_shapes = self.export_config.input_shapes
+        input_shapes = self.config.export.input_shapes
 
-        half_precision = int(self.trainer.precision) == 16
+        half_precision = "16" in self.trainer.precision
 
-        for export_type in self.export_config.types:
-            if export_type == "torchscript":
-                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
+        model_json, export_paths = export_model(
+            config=self.config,
+            model=model,
+            export_folder=self.export_folder,
+            half_precision=half_precision,
+            input_shapes=input_shapes,
+            idx_to_class={0: "good", 1: "defect"},
+        )
 
-                _, input_shapes = out
+        if len(export_paths) == 0:
+            return
 
-        model_json = {
-            "input_size": input_shapes,
-            "classes": {0: "good", 1: "defect"},
-            "mean": list(self.config.transforms.mean),
-            "std": list(self.config.transforms.std),
-            "image_threshold": np.round(self.module.image_threshold.value.item(), 3),
-            "pixel_threshold": np.round(self.module.pixel_threshold.value.item(), 3),
-        }
+        model_json["image_threshold"] = np.round(self.module.image_threshold.value.item(), 3)
+        model_json["pixel_threshold"] = np.round(self.module.pixel_threshold.value.item(), 3)
+        model_json["anomaly_method"] = self.config.model.model.name
 
         with open(os.path.join(self.export_folder, "model.json"), "w") as f:
             json.dump(model_json, f)
@@ -220,10 +201,12 @@ def _generate_report(self) -> None:
             if any(
                 isinstance(x, MinMaxNormalizationCallback) for x in self.trainer.callbacks  # type: ignore[attr-defined]
             )
-            else self.module.image_metrics.F1Score.threshold
+            else self.module.image_metrics.F1Score.threshold  # type: ignore[union-attr]
         )
 
-        plot_cumulative_histogram(good_scores, defect_scores, threshold.item(), self.report_path)
+        plot_cumulative_histogram(
+            good_scores, defect_scores, threshold.item(), self.report_path  # type: ignore[arg-type, operator]
+        )
 
         _, pd_cm, _ = get_results(np.array(gt_labels), np.array(pred_labels), idx_to_class)
         np_cm = np.array(pd_cm)
@@ -333,7 +316,7 @@ def test(self) -> None:
         self.datamodule.setup(stage="test")
         test_dataloader = self.datamodule.test_dataloader()
 
-        optimal_f1 = OptimalF1(num_classes=None, pos_label=1)
+        optimal_f1 = OptimalF1(num_classes=None, pos_label=1)  # type: ignore[arg-type]
 
         anomaly_scores = []
         anomaly_maps = []
@@ -346,7 +329,11 @@ def test(self) -> None:
                 batch_labels = batch_item["label"]
                 image_labels.extend(batch_labels.tolist())
                 image_paths.extend(batch_item["image_path"])
-                anomaly_map, anomaly_score = self.deployment_model(batch_images.to(self.device))
+                if self.model_data.get("anomaly_method") == "efficientad":
+                    model_output = self.deployment_model(batch_images.to(self.device), None)
+                else:
+                    model_output = self.deployment_model(batch_images.to(self.device))
+                anomaly_map, anomaly_score = model_output[0], model_output[1]
                 anomaly_map = anomaly_map.cpu()
                 anomaly_score = anomaly_score.cpu()
                 known_labels = torch.where(batch_labels != -1)[0]
diff --git a/quadra/tasks/base.py b/quadra/tasks/base.py
index 5ea6f820..cc77af11 100644
--- a/quadra/tasks/base.py
+++ b/quadra/tasks/base.py
@@ -11,13 +11,11 @@
 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.models.evaluation import BaseEvaluationModel
 from quadra.utils import utils
 from quadra.utils.export import import_deployment_model
 
@@ -30,16 +28,10 @@ class Task(Generic[DataModuleT]):
 
     Args:
         config: The experiment configuration.
-        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_config: Optional[DictConfig] = None):
+    def __init__(self, config: DictConfig):
         self.config = config
-        self.export_config = export_config
         self.export_folder: str = "deployment_model"
         self._datamodule: DataModuleT
         self.metadata: Dict[str, Any]
@@ -92,7 +84,7 @@ def execute(self) -> None:
         self.prepare()
         self.train()
         self.test()
-        if self.export_config is not None and len(self.export_config.types) > 0:
+        if self.config.export is not None and len(self.config.export.types) > 0:
             self.export()
         self.generate_report()
         self.finalize()
@@ -106,11 +98,6 @@ 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_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__(
@@ -119,10 +106,8 @@ def __init__(
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
         report: bool = False,
-        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config, export_config=export_config)
-        self.config = config
+        super().__init__(config=config)
         self.checkpoint_path = checkpoint_path
         self.run_test = run_test
         self.report = report
@@ -308,7 +293,7 @@ def execute(self) -> None:
         self.train()
         if self.run_test:
             self.test()
-        if self.export_config is not None and len(self.export_config.types) > 0:
+        if self.config.export is not None and len(self.config.export.types) > 0:
             self.export()
         if self.report:
             self.generate_report()
@@ -351,21 +336,23 @@ def __init__(
         self.config = config
         self.model_data: Dict[str, Any]
         self.model_path = model_path
-        self._deployment_model: Union[RecursiveScriptModule, Module]
+        self._deployment_model: BaseEvaluationModel
         self.deployment_model_type: str
         self.model_info_filename = "model.json"
         self.report_path = ""
         self.metadata = {"report_files": []}
 
     @property
-    def deployment_model(self) -> Union[RecursiveScriptModule, nn.Module]:
+    def deployment_model(self) -> BaseEvaluationModel:
         """Deployment model."""
         return self._deployment_model
 
     @deployment_model.setter
     def deployment_model(self, model_path: str):
         """Set the deployment model."""
-        self._deployment_model, self.deployment_model_type = import_deployment_model(model_path, self.device)
+        self._deployment_model = import_deployment_model(
+            model_path=model_path, device=self.device, inference_config=self.config.inference
+        )
 
     def prepare(self) -> None:
         """Prepare the evaluation."""
diff --git a/quadra/tasks/classification.py b/quadra/tasks/classification.py
index cf6d7144..6255876b 100644
--- a/quadra/tasks/classification.py
+++ b/quadra/tasks/classification.py
@@ -31,12 +31,13 @@
 from quadra.datasets.classification import ImageClassificationListDataset
 from quadra.models.base import ModelSignatureWrapper
 from quadra.models.classification import BaseNetworkBuilder
+from quadra.models.evaluation import BaseEvaluationModel, TorchEvaluationModel
 from quadra.modules.classification import ClassificationModule
 from quadra.tasks.base import Evaluation, LightningTask, Task
 from quadra.trainers.classification import SklearnClassificationTrainer
 from quadra.utils import utils
 from quadra.utils.classification import get_results, save_classification_result
-from quadra.utils.export import export_pytorch_model, export_torchscript_model, import_deployment_model
+from quadra.utils.export import export_model, import_deployment_model
 from quadra.utils.models import get_feature, is_vision_transformer
 from quadra.utils.vit_explainability import VitAttentionGradRollout
 
@@ -59,11 +60,6 @@ class Classification(Generic[ClassificationDataModuleT], LightningTask[Classific
     Args:
         config: The experiment configuration
         output: The otuput configuration.
-        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.
@@ -79,7 +75,6 @@ def __init__(
         output: DictConfig,
         checkpoint_path: Optional[str] = None,
         lr_multiplier: Optional[float] = None,
-        export_config: Optional[DictConfig] = None,
         gradcam: bool = False,
         report: bool = False,
         run_test: bool = False,
@@ -89,7 +84,6 @@ def __init__(
             checkpoint_path=checkpoint_path,
             run_test=run_test,
             report=report,
-            export_config=export_config,
         )
         self.output = output
         self.gradcam = gradcam
@@ -249,10 +243,6 @@ def test(self) -> None:
 
     def export(self) -> None:
         """Generate deployment models for the task."""
-        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(
                 "No `class_to_idx` found in the datamodule, class information will not be saved in the model.json"
@@ -262,55 +252,36 @@ def export(self) -> None:
             idx_to_class = {v: k for k, v in self.datamodule.class_to_idx.items()}
 
         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)
-
-        module = self.module.load_from_checkpoint(
-            best_model_path,
-            model=self.module.model,
-            optimizer=self.optimizer,
-            lr_scheduler=self.scheduler,
-            criterion=self.module.criterion,
-            gradcam=False,
-        )
-
-        input_shapes = self.export_config.input_shapes
+            log.warning("No checkpoint callback found in the trainer, exporting the last model weights")
+        else:
+            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)
 
-        # TODO: This breaks with bf16 precision!!!
-        half_precision = int(self.trainer.precision) == 16
+            module = self.module.load_from_checkpoint(
+                best_model_path,
+                model=self.module.model,
+                optimizer=self.optimizer,
+                lr_scheduler=self.scheduler,
+                criterion=self.module.criterion,
+                gradcam=False,
+            )
 
-        for export_type in self.export_config.types:
-            if export_type == "torchscript":
-                out = export_torchscript_model(
-                    model=module.model,
-                    input_shapes=input_shapes,
-                    output_path=self.export_folder,
-                    half_precision=half_precision,
-                )
+        input_shapes = self.config.export.input_shapes
 
-                if out is None:
-                    log.warning("Skipping torchscript export since the model is not supported")
-                    continue
+        # TODO: What happens if we have 64 precision?
+        half_precision = "16" in self.trainer.precision
 
-                _, input_shapes = out
-            elif export_type == "pytorch":
-                export_pytorch_model(
-                    model=module.model,
-                    output_path=self.export_folder,
-                )
-                with open(os.path.join(self.export_folder, "model_config.yaml"), "w") as f:
-                    OmegaConf.save(self.config.model, f, resolve=True)
-            else:
-                log.warning("Export type: %s not implemented", export_type)
+        self.model_json, export_paths = export_model(
+            config=self.config,
+            model=module.model,
+            export_folder=self.export_folder,
+            half_precision=half_precision,
+            input_shapes=input_shapes,
+            idx_to_class=idx_to_class,
+        )
 
-        self.model_json = {
-            "input_size": input_shapes,
-            "classes": idx_to_class,
-            "mean": list(self.config.transforms.mean),
-            "std": list(self.config.transforms.std),
-        }
+        if len(export_paths) == 0:
+            return
 
         with open(os.path.join(self.export_folder, self.deploy_info_file), "w") as f:
             json.dump(self.model_json, f)
@@ -442,11 +413,6 @@ 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_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__(
@@ -454,13 +420,12 @@ def __init__(
         config: DictConfig,
         output: DictConfig,
         device: str,
-        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_config=export_config)
+        super().__init__(config=config)
 
         self._device = device
         self.output = output
-        self._backbone: nn.Module
+        self._backbone: ModelSignatureWrapper
         self._trainer: SklearnClassificationTrainer
         self._model: ClassifierMixin
         self.metadata: Dict[str, Any] = {
@@ -505,7 +470,7 @@ def model(self, model_config: DictConfig):
         self._model = hydra.utils.instantiate(model_config)
 
     @property
-    def backbone(self) -> nn.Module:
+    def backbone(self) -> ModelSignatureWrapper:
         """Backbone: The backbone."""
         return self._backbone
 
@@ -656,48 +621,29 @@ def test_full_data(self) -> None:
 
     def export(self) -> None:
         """Generate deployment model for the task."""
-        if self.export_config is None or len(self.export_config.types) == 0:
+        if self.config.export is None or len(self.config.export.types) == 0:
             log.info("No export type specified skipping export")
             return
 
-        input_shapes = self.export_config.input_shapes
+        input_shapes = self.config.export.input_shapes
 
-        for export_type in self.export_config.types:
-            if export_type == "torchscript":
-                out = export_torchscript_model(
-                    model=self.backbone,
-                    input_shapes=input_shapes,
-                    output_path=self.export_folder,
-                    half_precision=False,
-                )
+        idx_to_class = {v: k for k, v in self.datamodule.full_dataset.class_to_idx.items()}
 
-                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
-                with open(os.path.join(self.export_folder, "backbone_config.yaml"), "w") as f:
-                    OmegaConf.save(self.config.backbone, f, resolve=True)
-                log.info("backbone_config.yaml saved (export type 'pytorch')")
-            else:
-                log.warning("Export type: %s not implemented", export_type)
+        model_json, export_paths = export_model(
+            config=self.config,
+            model=self.backbone,
+            export_folder=self.export_folder,
+            half_precision=False,
+            input_shapes=input_shapes,
+            idx_to_class=idx_to_class,
+            pytorch_model_type="backbone",
+        )
 
         dump(self.model, os.path.join(self.export_folder, "classifier.joblib"))
 
-        idx_to_class = {v: k for k, v in self.datamodule.full_dataset.class_to_idx.items()}
-
-        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(model_json, f)
+        if len(export_paths) > 0:
+            with open(os.path.join(self.export_folder, self.deploy_info_file), "w") as f:
+                json.dump(model_json, f)
 
     def generate_report(self) -> None:
         """Generate report for the task."""
@@ -742,7 +688,7 @@ def execute(self) -> None:
         if self.output.report:
             self.generate_report()
         self.train_full_data()
-        if self.export_config is not None and len(self.export_config.types) > 0:
+        if self.config.export is not None and len(self.config.export.types) > 0:
             self.export()
         if self.output.test_full_data:
             self.test_full_data()
@@ -758,6 +704,7 @@ class SklearnTestClassification(Evaluation[SklearnClassificationDataModuleT]):
         model_path: path to trained model generated from SklearnClassification task.
         device: the device where to run the model (cuda or cpu)
         gradcam: Whether to compute gradcams
+        **kwargs: Additional arguments to pass to the task
     """
 
     def __init__(
@@ -769,10 +716,10 @@ def __init__(
         gradcam: bool = False,
         **kwargs: Any,
     ):
-        super().__init__(config=config, model_path=model_path, device=device)
+        super().__init__(config=config, model_path=model_path, device=device, **kwargs)
         self.gradcam = gradcam
         self.output = output
-        self._backbone: nn.Module
+        self._backbone: BaseEvaluationModel
         self._classifier: ClassifierMixin
         self.class_to_idx: Dict[str, int]
         self.idx_to_class: Dict[int, str]
@@ -833,32 +780,38 @@ def classifier(self, classifier_path: str) -> None:
         self._classifier = load(classifier_path)
 
     @property
-    def backbone(self) -> nn.Module:
+    def backbone(self) -> BaseEvaluationModel:
         """Backbone: The backbone."""
         return self._backbone
 
     @backbone.setter
     def backbone(self, model_path: str) -> None:
         """Load backbone."""
-        file_extension = os.path.splitext(os.path.basename(model_path))[1]
-        if file_extension == ".yaml":
-            log.info("Model path points to '.yaml' file")
-            backbone_config_path = os.path.join(Path(model_path).parent, "backbone_config.yaml")
+        file_extension = os.path.splitext(model_path)[1]
+
+        model_architecture = None
+        if file_extension == ".pth":
+            backbone_config_path = os.path.join(Path(model_path).parent, "model_config.yaml")
             log.info("Loading backbone from config")
             backbone_config = OmegaConf.load(backbone_config_path)
 
             if backbone_config.metadata.get("checkpoint"):
                 log.info("Loading backbone from <%s>", backbone_config.metadata.checkpoint)
-                self._backbone = torch.load(backbone_config.metadata.checkpoint)
+                model_architecture = torch.load(backbone_config.metadata.checkpoint)
             else:
                 log.info("Loading backbone from <%s>", backbone_config.model["_target_"])
-                self._backbone = hydra.utils.instantiate(backbone_config.model)
-            self._backbone.eval()
-            self._backbone = self._backbone.to(self.device)
-        else:
-            log.info("Importing trained model")
-            self._backbone, model_type = import_deployment_model(model_path=model_path, device=self.device)
-            log.info("Imported %s model", model_type)
+                model_architecture = hydra.utils.instantiate(backbone_config.model)
+
+        self._backbone = import_deployment_model(
+            model_path=model_path,
+            device=self.device,
+            inference_config=self.config.inference,
+            model_architecture=model_architecture,
+        )
+
+        if self.gradcam and not isinstance(self._backbone, TorchEvaluationModel):
+            log.warning("Gradcam is supported only for pytorch models. Skipping gradcam")
+            self.gradcam = False
 
     @property
     def trainer(self) -> SklearnClassificationTrainer:
@@ -947,59 +900,62 @@ def __init__(
         self.gradcam = gradcam
         self.cam: GradCAM
 
-    @property
-    def model(self) -> nn.Module:
-        return self._model
-
-    @model.setter
-    def model(self, model_config: DictConfig) -> None:
-        self.pre_classifier = model_config  # type: ignore[assignment]
-        self.classifier = model_config  # type: ignore[assignment]
+    def get_torch_model(self, model_config: DictConfig) -> nn.Module:
+        """Instantiate the torch model from the config."""
+        pre_classifier = self.get_pre_classifier(model_config)
+        classifier = self.get_classifier(model_config)
         log.info("Instantiating backbone <%s>", model_config.model["_target_"])
-        self._model: nn.Module = hydra.utils.instantiate(
-            model_config.model, classifier=self.classifier, pre_classifier=self.pre_classifier, _convert_="partial"
-        )
 
-    @property
-    def pre_classifier(self) -> nn.Module:
-        return self._pre_classifier
+        return hydra.utils.instantiate(
+            model_config.model, classifier=classifier, pre_classifier=pre_classifier, _convert_="partial"
+        )
 
-    @pre_classifier.setter
-    def pre_classifier(self, model_config: DictConfig) -> None:
+    def get_pre_classifier(self, model_config: DictConfig) -> nn.Module:
+        """Instantiate the pre-classifier from the config."""
         if "pre_classifier" in model_config and model_config.pre_classifier is not None:
             log.info("Instantiating pre_classifier <%s>", model_config.pre_classifier["_target_"])
-            self._pre_classifier = hydra.utils.instantiate(model_config.pre_classifier, _convert_="partial")
+            pre_classifier = hydra.utils.instantiate(model_config.pre_classifier, _convert_="partial")
         else:
             log.info("No pre-classifier found in config: instantiate a torch.nn.Identity instead")
-            self._pre_classifier = nn.Identity()
+            pre_classifier = nn.Identity()
 
-    @property
-    def classifier(self) -> nn.Module:
-        return self._classifier
+        return pre_classifier
 
-    @classifier.setter
-    def classifier(self, model_config: DictConfig) -> None:
+    def get_classifier(self, model_config: DictConfig) -> nn.Module:
+        """Instantiate the classifier from the config."""
         if "classifier" in model_config:
             log.info("Instantiating classifier <%s>", model_config.classifier["_target_"])
-            self._classifier = hydra.utils.instantiate(model_config.classifier, _convert_="partial")
-        else:
-            raise ValueError("A `classifier` definition must be specified in the config")
+            return hydra.utils.instantiate(model_config.classifier, _convert_="partial")
+
+        raise ValueError("A `classifier` definition must be specified in the config")
 
     @property
-    def deployment_model(self):
+    def deployment_model(self) -> BaseEvaluationModel:
         """Deployment model."""
         return self._deployment_model
 
     @deployment_model.setter
     def deployment_model(self, model_path: str):
         """Set the deployment model."""
-        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(
-            model_path, self.device, self.model
+        file_extension = os.path.splitext(model_path)[1]
+
+        model_architecture = None
+        if file_extension == ".pth":
+            model_config = OmegaConf.load(os.path.join(Path(model_path).parent, "model_config.yaml"))
+
+            if not isinstance(model_config, DictConfig):
+                raise ValueError(f"The model config must be a DictConfig, got {type(model_config)}")
+
+            model_architecture = self.get_torch_model(model_config)
+
+        self._deployment_model = import_deployment_model(
+            model_path=model_path,
+            device=self.device,
+            inference_config=self.config.inference,
+            model_architecture=model_architecture,
         )
-        if self.gradcam and self.deployment_model_type != "torch":
+
+        if self.gradcam and not isinstance(self.deployment_model, TorchEvaluationModel):
             log.warning("To compute gradcams you need to provide the path to an exported .pth state_dict file")
             self.gradcam = False
 
@@ -1011,19 +967,26 @@ def prepare(self) -> None:
 
     def prepare_gradcam(self) -> None:
         """Initializing gradcam for the predictions."""
-        if isinstance(self.deployment_model.features_extractor, timm.models.resnet.ResNet):
+        if not hasattr(self.deployment_model.model, "features_extractor"):
+            log.warning("Gradcam not implemented for this backbone, it will not be computed")
+            self.gradcam = False
+            return
+
+        if isinstance(self.deployment_model.model.features_extractor, timm.models.resnet.ResNet):
             target_layers = [
-                cast(BaseNetworkBuilder, self.deployment_model).features_extractor.layer4[-1]  # type: ignore[index]
+                cast(BaseNetworkBuilder, self.deployment_model.model).features_extractor.layer4[
+                    -1
+                ]  # type: ignore[index]
             ]
             self.cam = GradCAM(
-                model=self.deployment_model,
+                model=self.deployment_model.model,
                 target_layers=target_layers,
                 use_cuda=(self.device != "cpu"),
             )
-            for p in self.deployment_model.features_extractor.layer4[-1].parameters():
+            for p in self.deployment_model.model.features_extractor.layer4[-1].parameters():
                 p.requires_grad = True
-        elif is_vision_transformer(cast(BaseNetworkBuilder, self.deployment_model).features_extractor):
-            self.grad_rollout = VitAttentionGradRollout(self.deployment_model)
+        elif is_vision_transformer(cast(BaseNetworkBuilder, self.deployment_model.model).features_extractor):
+            self.grad_rollout = VitAttentionGradRollout(cast(nn.Module, self.deployment_model.model))
         else:
             log.warning("Gradcam not implemented for this backbone, it will not be computed")
             self.gradcam = False
@@ -1037,6 +1000,7 @@ def test(self) -> None:
         test_dataloader = self.datamodule.test_dataloader()
 
         image_labels = []
+        probabilities = []
         predicted_classes = []
         grayscale_cams_list = []
 
@@ -1047,19 +1011,28 @@ def test(self) -> None:
             for batch_item in tqdm(test_dataloader):
                 im, target = batch_item
                 im = im.to(self.device).detach()
-                outputs = self.deployment_model(im).detach()
+
+                if self.gradcam:
+                    # When gradcam is used we need to remove gradients
+                    outputs = self.deployment_model(im).detach()
+                else:
+                    outputs = self.deployment_model(im)
+
                 probs = torch.softmax(outputs, dim=1)
                 preds = torch.max(probs, dim=1).indices
 
+                probabilities.append(probs.tolist())
                 predicted_classes.append(preds.tolist())
                 image_labels.extend(target.tolist())
-                if self.gradcam:
+                if self.gradcam and hasattr(self.deployment_model.model, "features_extractor"):
                     with torch.inference_mode(False):
                         im = im.clone()
-                        if isinstance(self.deployment_model.features_extractor, timm.models.resnet.ResNet):
+                        if isinstance(self.deployment_model.model.features_extractor, timm.models.resnet.ResNet):
                             grayscale_cam = self.cam(input_tensor=im, targets=None)
                             grayscale_cams_list.append(torch.from_numpy(grayscale_cam))
-                        elif is_vision_transformer(cast(BaseNetworkBuilder, self.deployment_model).features_extractor):
+                        elif is_vision_transformer(
+                            cast(BaseNetworkBuilder, self.deployment_model.model).features_extractor
+                        ):
                             grayscale_cam_low_res = self.grad_rollout(input_tensor=im, targets_list=preds.tolist())
                             orig_shape = grayscale_cam_low_res.shape
                             new_shape = (orig_shape[0], im.shape[2], im.shape[3])
@@ -1072,6 +1045,7 @@ def test(self) -> None:
             grayscale_cams = torch.cat(grayscale_cams_list, dim=0)
 
         predicted_classes = [item for sublist in predicted_classes for item in sublist]
+        probabilities = [item for sublist in probabilities for item in sublist]
         if self.datamodule.class_to_idx is not None:
             idx_to_class = {v: k for k, v in self.datamodule.class_to_idx.items()}
 
@@ -1086,6 +1060,7 @@ def test(self) -> None:
         self.metadata["test_confusion_matrix"] = pd_cm
         self.metadata["test_accuracy"] = test_accuracy
         self.metadata["test_results"] = predicted_classes
+        self.metadata["probabilities"] = probabilities
         self.metadata["test_labels"] = image_labels
         self.metadata["grayscale_cams"] = grayscale_cams
 
diff --git a/quadra/tasks/patch.py b/quadra/tasks/patch.py
index afb0def6..c0d6d5b7 100644
--- a/quadra/tasks/patch.py
+++ b/quadra/tasks/patch.py
@@ -1,7 +1,7 @@
 import json
 import os
 from pathlib import Path
-from typing import Any, Dict, List, Optional, cast
+from typing import Any, Dict, List, cast
 
 import hydra
 import torch
@@ -12,10 +12,11 @@
 from quadra.datamodules import PatchSklearnClassificationDataModule
 from quadra.datasets.patch import PatchSklearnClassificationTrainDataset
 from quadra.models.base import ModelSignatureWrapper
+from quadra.models.evaluation import BaseEvaluationModel
 from quadra.tasks.base import Evaluation, Task
 from quadra.trainers.classification import SklearnClassificationTrainer
 from quadra.utils import utils
-from quadra.utils.export import export_torchscript_model, import_deployment_model
+from quadra.utils.export import export_model, import_deployment_model
 from quadra.utils.patch import RleEncoder, compute_patch_metrics, save_classification_result
 from quadra.utils.patch.dataset import PatchDatasetFileFormat
 
@@ -29,11 +30,6 @@ 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_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__(
@@ -41,9 +37,8 @@ def __init__(
         config: DictConfig,
         output: DictConfig,
         device: str,
-        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_config=export_config)
+        super().__init__(config=config)
         self.device: str = device
         self.output: DictConfig = output
         self.return_polygon: bool = True
@@ -207,65 +202,53 @@ def generate_report(self) -> None:
 
     def export(self) -> None:
         """Generate deployment model for the task."""
-        if self.export_config is None or len(self.export_config.types) == 0:
-            log.info("No export type specified skipping export")
-            return
-
-        os.makedirs(self.export_folder, exist_ok=True)
+        input_shapes = self.config.export.input_shapes
 
-        input_shapes = self.export_config.input_shapes
+        idx_to_class = {v: k for k, v in self.datamodule.class_to_idx.items()}
 
-        for export_type in self.export_config.types:
-            if export_type == "torchscript":
-                out = export_torchscript_model(
-                    model=self.backbone,
-                    input_shapes=input_shapes,
-                    output_path=self.export_folder,
-                    half_precision=False,
-                )
+        model_json, export_paths = export_model(
+            config=self.config,
+            model=self.backbone,
+            export_folder=self.export_folder,
+            half_precision=False,
+            input_shapes=input_shapes,
+            idx_to_class=idx_to_class,
+            pytorch_model_type="backbone",
+        )
 
-                if out is None:
-                    log.warning("Skipping torchscript export since the model is not supported")
-                    continue
+        if len(export_paths) > 0:
+            dataset_info = self.datamodule.info
+
+            horizontal_patches = dataset_info.patch_number[1] if dataset_info.patch_number is not None else None
+            vertical_patches = dataset_info.patch_number[0] if dataset_info.patch_number is not None else None
+            patch_height = dataset_info.patch_size[0] if dataset_info.patch_size is not None else None
+            patch_width = dataset_info.patch_size[1] if dataset_info.patch_size is not None else None
+            overlap = dataset_info.overlap
+
+            model_json.update(
+                {
+                    "horizontal_patches": horizontal_patches,
+                    "vertical_patches": vertical_patches,
+                    "patch_height": patch_height,
+                    "patch_width": patch_width,
+                    "overlap": overlap,
+                    "reconstruction_method": self.output.reconstruction_method,
+                    "class_to_skip": self.datamodule.class_to_skip_training,
+                }
+            )
 
-                _, input_shapes = out
+            with open(os.path.join(self.export_folder, "model.json"), "w") as f:
+                json.dump(model_json, f, cls=utils.HydraEncoder)
 
         dump(self.model, os.path.join(self.export_folder, "classifier.joblib"))
 
-        dataset_info = self.datamodule.info
-
-        horizontal_patches = dataset_info.patch_number[1] if dataset_info.patch_number is not None else None
-        vertical_patches = dataset_info.patch_number[0] if dataset_info.patch_number is not None else None
-        patch_height = dataset_info.patch_size[0] if dataset_info.patch_size is not None else None
-        patch_width = dataset_info.patch_size[1] if dataset_info.patch_size is not None else None
-        overlap = dataset_info.overlap
-
-        idx_to_class = {v: k for k, v in self.datamodule.class_to_idx.items()}
-
-        model_json = {
-            "input_size": input_shapes,
-            "classes": idx_to_class,
-            "mean": self.config.transforms.mean,
-            "std": self.config.transforms.std,
-            "horizontal_patches": horizontal_patches,
-            "vertical_patches": vertical_patches,
-            "patch_height": patch_height,
-            "patch_width": patch_width,
-            "overlap": overlap,
-            "reconstruction_method": self.output.reconstruction_method,
-            "class_to_skip": self.datamodule.class_to_skip_training,
-        }
-
-        with open(os.path.join(self.export_folder, "model.json"), "w") as f:
-            json.dump(model_json, f, cls=utils.HydraEncoder)
-
     def execute(self) -> None:
         """Execute the experiment and all the steps."""
         self.prepare()
         self.train()
         if self.output.report:
             self.generate_report()
-        if self.export_config is not None and len(self.export_config.types) > 0:
+        if self.config.export is not None and len(self.config.export.types) > 0:
             self.export()
         self.finalize()
 
@@ -289,7 +272,7 @@ def __init__(
     ):
         super().__init__(config=config, model_path=model_path, device=device)
         self.output = output
-        self._backbone: torch.nn.Module
+        self._backbone: BaseEvaluationModel
         self._classifier: ClassifierMixin
         self.class_to_idx: Dict[str, int]
         self.idx_to_class: Dict[int, str]
@@ -371,32 +354,34 @@ def classifier(self, classifier_path: str) -> None:
         self._classifier = load(classifier_path)
 
     @property
-    def backbone(self) -> torch.nn.Module:
+    def backbone(self) -> BaseEvaluationModel:
         """Backbone: The backbone."""
         return self._backbone
 
     @backbone.setter
     def backbone(self, model_path: str) -> None:
         """Load backbone."""
-        file_extension = os.path.splitext(os.path.basename(model_path))[1]
-        if file_extension == ".yaml":
-            log.info("Model path points to '.yaml' file")
-            backbone_config_path = os.path.join(Path(model_path).parent, "backbone_config.yaml")
+        file_extension = os.path.splitext(model_path)[1]
+
+        model_architecture = None
+        if file_extension == ".pth":
+            backbone_config_path = os.path.join(Path(model_path).parent, "model_config.yaml")
             log.info("Loading backbone from config")
             backbone_config = OmegaConf.load(backbone_config_path)
 
             if backbone_config.metadata.get("checkpoint"):
                 log.info("Loading backbone from <%s>", backbone_config.metadata.checkpoint)
-                self._backbone = torch.load(backbone_config.metadata.checkpoint)
+                model_architecture = torch.load(backbone_config.metadata.checkpoint)
             else:
                 log.info("Loading backbone from <%s>", backbone_config.model["_target_"])
-                self._backbone = hydra.utils.instantiate(backbone_config.model)
-            self._backbone.eval()
-            self._backbone = self._backbone.to(self.device)
-        else:
-            log.info("Importing trained model")
-            self._backbone, model_type = import_deployment_model(model_path=model_path, device=self.device)
-            log.info("Imported %s model", model_type)
+                model_architecture = hydra.utils.instantiate(backbone_config.model)
+
+        self._backbone = import_deployment_model(
+            model_path=model_path,
+            device=self.device,
+            inference_config=self.config.inference,
+            model_architecture=model_architecture,
+        )
 
     @property
     def trainer(self) -> SklearnClassificationTrainer:
diff --git a/quadra/tasks/segmentation.py b/quadra/tasks/segmentation.py
index 11de316c..a9ff8dc1 100644
--- a/quadra/tasks/segmentation.py
+++ b/quadra/tasks/segmentation.py
@@ -12,11 +12,12 @@
 from quadra.callbacks.mlflow import get_mlflow_logger
 from quadra.datamodules import SegmentationDataModule, SegmentationMulticlassDataModule
 from quadra.models.base import ModelSignatureWrapper
+from quadra.models.evaluation import BaseEvaluationModel
 from quadra.modules.base import SegmentationModel
 from quadra.tasks.base import Evaluation, LightningTask
 from quadra.utils import utils
 from quadra.utils.evaluation import create_mask_report
-from quadra.utils.export import export_torchscript_model
+from quadra.utils.export import export_model
 
 log = utils.get_logger(__name__)
 
@@ -35,11 +36,6 @@ 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_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 +46,12 @@ def __init__(
         run_test: bool = False,
         evaluate: Optional[DictConfig] = None,
         report: bool = False,
-        export_config: Optional[DictConfig] = None,
     ):
         super().__init__(
             config=config,
             checkpoint_path=checkpoint_path,
             run_test=run_test,
             report=report,
-            export_config=export_config,
         )
         self.evaluate = evaluate
         self.num_viz_samples = num_viz_samples
@@ -65,18 +59,18 @@ def __init__(
         self.exported_model_path: Optional[str] = None
         if self.evaluate and any(self.evaluate.values()):
             if (
-                self.export_config is None
-                or len(self.export_config.types) == 0
-                or "torchscript" not in self.export_config.types
+                self.config.export is None
+                or len(self.config.export.types) == 0
+                or "torchscript" not in self.config.export.types
             ):
                 log.info(
                     "Evaluation is enabled, but training does not export a deployment model. Automatically export the "
                     "model as torchscript."
                 )
-                if self.export_config is None:
-                    self.export_config = DictConfig({"types": ["torchscript"]})
+                if self.config.export is None:
+                    self.config.export = DictConfig({"types": ["torchscript"]})
                 else:
-                    self.export_config.types.append("torchscript")
+                    self.config.export.types.append("torchscript")
 
             if not self.report:
                 log.info("Evaluation is enabled, but reporting is disabled. Enabling reporting automatically.")
@@ -126,59 +120,51 @@ def export(self) -> None:
         log.info("Exporting model ready for deployment")
         # Get best model!
         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("Loaded best model from %s", best_model_path)
-
-        module = self.module.load_from_checkpoint(
-            best_model_path,
-            model=self.module.model,
-            loss_fun=None,
-            optimizer=self.module.optimizer,
-            lr_scheduler=self.module.schedulers,
-        )
+            log.warning("No checkpoint callback found in the trainer, exporting the last model weights")
+            module = self.module
+        else:
+            best_model_path = self.trainer.checkpoint_callback.best_model_path  # type: ignore[attr-defined]
+            log.info("Loaded best model from %s", best_model_path)
+
+            module = self.module.load_from_checkpoint(
+                best_model_path,
+                model=self.module.model,
+                loss_fun=None,
+                optimizer=self.module.optimizer,
+                lr_scheduler=self.module.schedulers,
+            )
 
         if "idx_to_class" not in self.config.datamodule:
             log.info("No idx_to_class key")
-            classes = {0: "good", 1: "bad"}
+            idx_to_class = {0: "good", 1: "bad"}  # TODO: Why is this the default value?
         else:
             log.info("idx_to_class is present")
-            classes = self.config.datamodule.idx_to_class
+            idx_to_class = self.config.datamodule.idx_to_class
 
-        if self.export_config is None:
+        if self.config.export 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."
             )
 
-        input_shapes = self.export_config.input_shapes
+        half_precision = "16" in self.trainer.precision
 
-        half_precision = self.trainer.precision == 16
+        input_shapes = self.config.export.input_shapes
 
-        for export_type in self.export_config.types:
-            if export_type == "torchscript":
-                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
+        model_json, export_paths = export_model(
+            config=self.config,
+            model=module.model,
+            export_folder=self.export_folder,
+            half_precision=half_precision,
+            input_shapes=input_shapes,
+            idx_to_class=idx_to_class,
+        )
 
-        if input_shapes is None:
-            log.warning("Not able to export the model in any format")
+        if len(export_paths) == 0:
+            return
 
-        model_json = {
-            "input_size": input_shapes,
-            "classes": classes,
-            "mean": self.config.transforms.mean,
-            "std": self.config.transforms.std,
-        }
+        # Pick one model for evaluation, it should be independent of the export type as the model is wrapped
+        self.exported_model_path = next(iter(export_paths.values()))
 
         with open(os.path.join(self.export_folder, "model.json"), "w") as f:
             json.dump(model_json, f, cls=utils.HydraEncoder)
@@ -272,7 +258,7 @@ def prepare(self) -> None:
 
     @torch.no_grad()
     def inference(
-        self, dataloader: DataLoader, deployment_model: torch.nn.Module, device: torch.device
+        self, dataloader: DataLoader, deployment_model: BaseEvaluationModel, device: torch.device
     ) -> Dict[str, torch.Tensor]:
         """Run inference on the dataloader and return the output.
 
diff --git a/quadra/tasks/ssl.py b/quadra/tasks/ssl.py
index 84cd3af8..0804154b 100644
--- a/quadra/tasks/ssl.py
+++ b/quadra/tasks/ssl.py
@@ -1,22 +1,22 @@
 import json
 import os
-from typing import Any, List, Optional, Tuple, Union, cast
+from typing import Any, List, Optional, Tuple, cast
 
 import hydra
 import torch
 from omegaconf import DictConfig, open_dict
 from pytorch_lightning import LightningModule
 from torch import nn
-from torch.jit._script import RecursiveScriptModule
 from torch.nn.functional import interpolate
 from torch.utils.tensorboard import SummaryWriter
 from tqdm import tqdm
 
 from quadra.callbacks.scheduler import WarmupInit
 from quadra.models.base import ModelSignatureWrapper
+from quadra.models.evaluation import BaseEvaluationModel
 from quadra.tasks.base import LightningTask, Task
 from quadra.utils import utils
-from quadra.utils.export import export_torchscript_model, import_deployment_model
+from quadra.utils.export import export_model, import_deployment_model
 
 log = utils.get_logger(__name__)
 
@@ -29,11 +29,6 @@ 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_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__(
@@ -42,14 +37,12 @@ def __init__(
         run_test: bool = False,
         report: bool = False,
         checkpoint_path: Optional[str] = None,
-        export_config: Optional[DictConfig] = None,
     ):
         super().__init__(
             config=config,
             checkpoint_path=checkpoint_path,
             run_test=run_test,
             report=report,
-            export_config=export_config,
         )
         self._backbone: nn.Module
         self._optimizer: torch.optim.Optimizer
@@ -100,37 +93,21 @@ def test(self) -> None:
 
     def export(self) -> None:
         """Deploy a model ready for production."""
-        if self.export_config is None or len(self.export_config.types) == 0:
-            log.info("No export type specified skipping export")
-            return
-
-        mean = self.config.transforms.mean
-        std = self.config.transforms.std
-        half_precision = int(self.trainer.precision) == 16
-
-        input_shapes = self.export_config.input_shapes
+        half_precision = "16" in self.trainer.precision
 
-        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,
-                )
+        input_shapes = self.config.export.input_shapes
 
-                if out is None:
-                    log.warning("Skipping torchscript export since the model is not supported")
-                    continue
-
-                _, input_shapes = out
+        model_json, export_paths = export_model(
+            config=self.config,
+            model=self.module.model,
+            export_folder=self.export_folder,
+            half_precision=half_precision,
+            input_shapes=input_shapes,
+            idx_to_class=None,
+        )
 
-        model_json = {
-            "input_size": input_shapes,
-            "classes": None,
-            "mean": mean,
-            "std": std,
-        }
+        if len(export_paths) == 0:
+            return
 
         with open(os.path.join(self.export_folder, "model.json"), "w") as f:
             json.dump(model_json, f, cls=utils.HydraEncoder)
@@ -145,11 +122,6 @@ class Simsiam(SSL):
             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__(
@@ -157,9 +129,8 @@ def __init__(
         config: DictConfig,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
-        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_config=export_config, checkpoint_path=checkpoint_path, run_test=run_test)
+        super().__init__(config=config, checkpoint_path=checkpoint_path, run_test=run_test)
         self.backbone: nn.Module
         self.projection_mlp: nn.Module
         self.prediction_mlp: nn.Module
@@ -220,11 +191,6 @@ class SimCLR(SSL):
             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__(
@@ -232,9 +198,8 @@ def __init__(
         config: DictConfig,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
-        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_config=export_config, checkpoint_path=checkpoint_path, run_test=run_test)
+        super().__init__(config=config, checkpoint_path=checkpoint_path, run_test=run_test)
         self.backbone: nn.Module
         self.projection_mlp: nn.Module
 
@@ -287,11 +252,6 @@ class Barlow(SimCLR):
             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__(
@@ -299,9 +259,8 @@ def __init__(
         config: DictConfig,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
-        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_config=export_config, checkpoint_path=checkpoint_path, run_test=run_test)
+        super().__init__(config=config, checkpoint_path=checkpoint_path, run_test=run_test)
 
     def prepare(self) -> None:
         """Prepare the experiment."""
@@ -330,11 +289,6 @@ class BYOL(SSL):
             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.
         **kwargs: Keyword arguments
     """
 
@@ -343,12 +297,10 @@ def __init__(
         config: DictConfig,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
-        export_config: Optional[DictConfig] = None,
         **kwargs: Any,
     ):
         super().__init__(
             config=config,
-            export_config=export_config,
             checkpoint_path=checkpoint_path,
             run_test=run_test,
             **kwargs,
@@ -420,11 +372,6 @@ class DINO(SSL):
             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__(
@@ -432,9 +379,8 @@ def __init__(
         config: DictConfig,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
-        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_config=export_config, checkpoint_path=checkpoint_path, run_test=run_test)
+        super().__init__(config=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
@@ -529,7 +475,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: Union[RecursiveScriptModule, nn.Module]
+        self._deployment_model: BaseEvaluationModel
         self.deployment_model_type: str
 
     @property
@@ -542,8 +488,7 @@ def device(self, device: str):
 
         if self.deployment_model is not None:
             # After prepare
-            if isinstance(self.deployment_model, RecursiveScriptModule):
-                self.deployment_model = self.deployment_model.to(self._device)
+            self.deployment_model = self.deployment_model.to(self._device)
 
     @property
     def deployment_model(self):
@@ -553,7 +498,9 @@ def deployment_model(self):
     @deployment_model.setter
     def deployment_model(self, model_path: str):
         """Set the deployment model."""
-        self._deployment_model, self.deployment_model_type = import_deployment_model(model_path, self.device)
+        self._deployment_model = import_deployment_model(
+            model_path=model_path, device=self.device, inference_config=self.config.inference
+        )
 
     def prepare(self) -> None:
         """Prepare the evaluation."""
diff --git a/quadra/utils/export.py b/quadra/utils/export.py
index c66040c4..47a7a0f5 100644
--- a/quadra/utils/export.py
+++ b/quadra/utils/export.py
@@ -1,18 +1,39 @@
 import os
-from typing import Any, List, Optional, Sequence, Tuple, Union, cast
+from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, TypeVar, Union, cast
 
 import torch
 from anomalib.models.cflow import CflowLightning
+from omegaconf import DictConfig, OmegaConf
 from torch import nn
-from torch.jit._script import RecursiveScriptModule
 
 from quadra.models.base import ModelSignatureWrapper
+from quadra.models.evaluation import (
+    BaseEvaluationModel,
+    ONNXEvaluationModel,
+    TorchEvaluationModel,
+    TorchscriptEvaluationModel,
+)
+from quadra.utils.logger import get_logger
 
-# TODO: Solve circular import as it is not possible to import get_logger right now
+try:
+    import onnx  # noqa
+    from onnxsim import simplify as onnx_simplify  # noqa
+
+    ONNX_AVAILABLE = True
+except ImportError:
+    ONNX_AVAILABLE = False
+
+log = get_logger(__name__)
+
+BaseDeploymentModelT = TypeVar("BaseDeploymentModelT", bound=BaseEvaluationModel)
 
 
 def generate_torch_inputs(
-    input_shapes: List[Any], device: str, half_precision: bool = False, dtype: torch.dtype = torch.float32
+    input_shapes: List[Any],
+    device: str,
+    half_precision: bool = False,
+    dtype: torch.dtype = torch.float32,
+    batch_size: int = 1,
 ) -> 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.
@@ -22,7 +43,7 @@ def generate_torch_inputs(
             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)
+        inp = torch.randn((batch_size, *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()}
@@ -33,7 +54,7 @@ def generate_torch_inputs(
             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)
+        inp = torch.randn((batch_size, *input_shapes), dtype=dtype, device=device)
 
     if half_precision:
         inp = inp.half()
@@ -41,6 +62,44 @@ def generate_torch_inputs(
     return inp
 
 
+def extract_torch_model_inputs(
+    model: Union[nn.Module, ModelSignatureWrapper],
+    input_shapes: Optional[List[Any]] = None,
+    half_precision: bool = False,
+    batch_size: int = 1,
+) -> Optional[Tuple[Union[List[Any], Tuple[Any, ...], torch.Tensor], List[Any]]]:
+    """Extract the input shapes for the given model and generate a list of torch tensors with the
+    given device and dtype.
+
+    Args:
+        model: Module or ModelSignatureWrapper
+        input_shapes: Inputs shapes
+        half_precision: If True, the model will be exported with half precision
+        batch_size: Batch size for the input shapes
+    """
+    if isinstance(model, ModelSignatureWrapper):
+        if input_shapes is None:
+            input_shapes = model.input_shapes
+
+    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 half_precision:
+        inp = generate_torch_inputs(
+            input_shapes=input_shapes, device="cuda:0", half_precision=True, dtype=torch.float16, batch_size=batch_size
+        )
+    else:
+        inp = generate_torch_inputs(
+            input_shapes=input_shapes, device="cpu", half_precision=False, dtype=torch.float32, batch_size=batch_size
+        )
+
+    return inp, input_shapes
+
+
+@torch.inference_mode()
 def export_torchscript_model(
     model: nn.Module,
     output_path: str,
@@ -61,43 +120,174 @@ def export_torchscript_model(
         If the model is exported successfully, the path to the model and the input shape are returned.
 
     """
+    if isinstance(model, CflowLightning):
+        log.warning("Exporting cflow model with torchscript is not supported yet.")
+        return None
+
+    model.eval()
+    if half_precision:
+        model.to("cuda:0")
+        model = model.half()
+    else:
+        model.cpu()
+
+    model_inputs = extract_torch_model_inputs(model, input_shapes, half_precision)
+
+    if model_inputs is None:
+        return None
+
     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."
-        # )
+    inp, input_shapes = model_inputs
+
+    try:
+        try:
+            model_jit = torch.jit.trace(model, inp)
+        except RuntimeError as e:
+            log.warning("Standard tracing failed with exception %s, attempting tracing with strict=False", e)
+            model_jit = torch.jit.trace(model, inp, strict=False)
+
+        os.makedirs(output_path, exist_ok=True)
+
+        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))
+
+        return os.path.join(os.getcwd(), model_path), input_shapes
+    except Exception as e:
+        log.debug("Failed to export torchscript model with exception: %s", e)
         return None
 
-    if isinstance(model, CflowLightning):
-        # log.warning("Exporting cflow model with torchscript is not supported yet.")
+
+@torch.inference_mode()
+def export_onnx_model(
+    model: nn.Module,
+    output_path: str,
+    onnx_config: DictConfig,
+    input_shapes: Optional[List[Any]] = None,
+    half_precision: bool = False,
+    model_name: str = "model.onnx",
+) -> Optional[Tuple[str, Any]]:
+    """Export a PyTorch model with ONNX.
+
+    Args:
+        model: PyTorch model to be exported
+        output_path: Path to save the model
+        input_shapes: Input shapes for tracing
+        onnx_config: ONNX export configuration
+        half_precision: If True, the model will be exported with half precision
+        model_name: Name of the exported model
+    """
+    if not ONNX_AVAILABLE:
+        log.warning("ONNX is not installed, can not export model in this format.")
+        log.warning("Please install ONNX capabilities for quadra with: pip install .[onnx]")
         return None
 
     model.eval()
     if half_precision:
-        # log.info("Jitting model with half precision on GPU")
         model.to("cuda:0")
         model = model.half()
-        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")
         model.cpu()
-        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)
+    if hasattr(onnx_config, "fixed_batch_size") and onnx_config.fixed_batch_size is not None:
+        batch_size = onnx_config.fixed_batch_size
+    else:
+        batch_size = 1
+
+    model_inputs = extract_torch_model_inputs(
+        model=model, input_shapes=input_shapes, half_precision=half_precision, batch_size=batch_size
+    )
+    if model_inputs is None:
+        return None
+
+    if isinstance(model, ModelSignatureWrapper):
+        model = model.instance
+
+    inp, input_shapes = model_inputs
 
     os.makedirs(output_path, exist_ok=True)
 
     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))
+    input_names = onnx_config.input_names if hasattr(onnx_config, "input_names") else None
+
+    if input_names is None:
+        input_names = []
+        for i, _ in enumerate(inp):
+            input_names.append(f"input_{i}")
+
+    output = [model(*inp)]
+    output_names = onnx_config.output_names if hasattr(onnx_config, "output_names") else None
+
+    if output_names is None:
+        output_names = []
+        for i, _ in enumerate(output):
+            output_names.append(f"output_{i}")
+
+    dynamic_axes = onnx_config.dynamic_axes if hasattr(onnx_config, "dynamic_axes") else None
+
+    if hasattr(onnx_config, "fixed_batch_size") and onnx_config.fixed_batch_size is not None:
+        dynamic_axes = None
+    else:
+        if dynamic_axes is None:
+            dynamic_axes = {}
+            for i, _ in enumerate(input_names):
+                dynamic_axes[input_names[i]] = {0: "batch_size"}
+
+            for i, _ in enumerate(output_names):
+                dynamic_axes[output_names[i]] = {0: "batch_size"}
+
+    onnx_config = cast(Dict[str, Any], OmegaConf.to_container(onnx_config, resolve=True))
+
+    onnx_config["input_names"] = input_names
+    onnx_config["output_names"] = output_names
+    onnx_config["dynamic_axes"] = dynamic_axes
+
+    simplify = onnx_config.pop("simplify", False)
+    _ = onnx_config.pop("fixed_batch_size", None)
+
+    if len(inp) == 1:
+        inp = inp[0]
+
+    if isinstance(inp, list):
+        inp = tuple(inp)  # onnx doesn't like lists representing tuples of inputs
+
+    if isinstance(inp, dict):
+        raise ValueError("ONNX export does not support model with dict inputs")
+
+    try:
+        torch.onnx.export(model=model, args=inp, f=model_path, **onnx_config)
+
+        onnx_model = onnx.load(model_path)
+        # Check if ONNX model is valid
+        onnx.checker.check_model(onnx_model)
+    except Exception as e:
+        log.debug("ONNX export failed with error: %s", e)
+        return None
+
+    log.info("ONNX model saved to %s", os.path.join(os.getcwd(), model_path))
+
+    if simplify:
+        log.info("Attempting to simplify ONNX model")
+        onnx_model = onnx.load(model_path)
+
+        try:
+            simplified_model, check = onnx_simplify(onnx_model)
+        except Exception as e:
+            log.debug("ONNX simplification failed with error: %s", e)
+            check = False
+
+        if not check:
+            log.warning("Something failed during model simplification, only original ONNX model will be exported")
+        else:
+            model_filename, model_extension = os.path.splitext(model_name)
+            model_name = f"{model_filename}_simplified{model_extension}"
+            model_path = os.path.join(output_path, model_name)
+            onnx.save(simplified_model, model_path)
+            log.info("Simplified ONNX model saved to %s", os.path.join(os.getcwd(), model_path))
 
     return os.path.join(os.getcwd(), model_path), input_shapes
 
@@ -114,47 +304,173 @@ def export_pytorch_model(model: nn.Module, output_path: str, model_name: str = "
         If the model is exported successfully, the path to the model is returned.
 
     """
+    if isinstance(model, ModelSignatureWrapper):
+        model = model.instance
+
     os.makedirs(output_path, exist_ok=True)
     model.eval()
     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)
 
 
-# TODO: Update signature when new models are added
+def export_model(
+    config: DictConfig,
+    model: Any,
+    export_folder: str,
+    half_precision: bool,
+    input_shapes: Optional[List[Any]] = None,
+    idx_to_class: Optional[Dict[int, str]] = None,
+    pytorch_model_type: Literal["backbone", "model"] = "model",
+) -> Tuple[Dict[str, Any], Dict[str, str]]:
+    """Generate deployment models for the task.
+
+    Args:
+        config: Experiment config
+        model: Model to be exported
+        export_folder: Path to save the exported model
+        half_precision: Whether to use half precision for the exported model
+        input_shapes: Input shapes for the exported model
+        idx_to_class: Mapping from class index to class name
+        pytorch_model_type: Type of the pytorch model config to be exported, if it's backbone on disk we will save the
+            config.backbone config, otherwise we will save the config.model
+
+    Returns:
+        If the model is exported successfully, return a dictionary containing information about the exported model and
+        a second dictionary containing the paths to the exported models. Otherwise, return two empty dictionaries.
+    """
+    if config.export is None or len(config.export.types) == 0:
+        log.info("No export type specified skipping export")
+        return {}, {}
+
+    os.makedirs(export_folder, exist_ok=True)
+
+    if input_shapes is None:
+        # Try to get input shapes from config
+        # If this is also None we will try to retrieve it from the ModelSignatureWrapper, if it fails we can't export
+        input_shapes = config.export.input_shapes
+
+    export_paths = {}
+
+    for export_type in config.export.types:
+        if export_type == "torchscript":
+            out = export_torchscript_model(
+                model=model,
+                input_shapes=input_shapes,
+                output_path=export_folder,
+                half_precision=half_precision,
+            )
+
+            if out is None:
+                log.warning("Torchscript export failed, enable debug logging for more details")
+                continue
+
+            export_path, input_shapes = out
+            export_paths[export_type] = export_path
+        elif export_type == "pytorch":
+            export_path = export_pytorch_model(
+                model=model,
+                output_path=export_folder,
+            )
+            export_paths[export_type] = export_path
+            with open(os.path.join(export_folder, "model_config.yaml"), "w") as f:
+                OmegaConf.save(getattr(config, pytorch_model_type), f, resolve=True)
+        elif export_type == "onnx":
+            if not hasattr(config.export, "onnx"):
+                log.warning("No onnx configuration found, skipping onnx export")
+                continue
+
+            out = export_onnx_model(
+                model=model,
+                output_path=export_folder,
+                onnx_config=config.export.onnx,
+                input_shapes=input_shapes,
+                half_precision=half_precision,
+            )
+
+            if out is None:
+                log.warning("ONNX export failed, enable debug logging for more details")
+                continue
+
+            export_path, input_shapes = out
+            export_paths[export_type] = export_path
+        else:
+            log.warning("Export type: %s not implemented", export_type)
+
+    if len(export_paths) == 0:
+        log.warning("No export type was successful, no model will be available for deployment")
+        return {}, export_paths
+
+    model_json = {
+        "input_size": input_shapes,
+        "classes": idx_to_class,
+        "mean": list(config.transforms.mean),
+        "std": list(config.transforms.std),
+    }
+
+    return model_json, export_paths
+
+
 def import_deployment_model(
-    model_path: str, device: str, model: Optional[nn.Module] = None
-) -> Tuple[Union[RecursiveScriptModule, nn.Module], str]:
+    model_path: str, inference_config: DictConfig, device: str, model_architecture: Optional[nn.Module] = None
+) -> BaseEvaluationModel:
     """Try to import a model for deployment, currently only supports torchscript .pt files and
     state dictionaries .pth files.
 
     Args:
         model_path: Path to the model
+        inference_config: Inference configuration, should contain keys for the different deployment models
         device: Device to load the model on
-        model: Pytorch model needed to load the parameter dictionary
+        model_architecture: Optional model architecture to use for loading a plain pytorch model
 
     Returns:
         A tuple containing the model and the model type
     """
+    log.info("Importing trained model")
 
     file_extension = os.path.splitext(os.path.basename(model_path))[1]
+    deployment_model: Optional[BaseEvaluationModel] = None
+
     if file_extension == ".pt":
-        model = cast(RecursiveScriptModule, torch.jit.load(model_path))
-        model.eval()
-        model.to(device)
+        deployment_model = TorchscriptEvaluationModel(config=inference_config.torchscript)
+    elif file_extension == ".pth":
+        if model_architecture is None:
+            raise ValueError("model_architecture must be specified when loading a .pth file")
 
-        return model, "torchscript"
-    if file_extension == ".pth":
-        if model is None:
-            raise ValueError("Model is not defined, can not load state_dict!")
+        deployment_model = TorchEvaluationModel(config=inference_config.pytorch, model_architecture=model_architecture)
+    elif file_extension == ".onnx":
+        deployment_model = ONNXEvaluationModel(config=inference_config.onnx)
 
-        model.load_state_dict(torch.load(model_path))
-        model.eval()
-        model.to(device)
+    if deployment_model is not None:
+        deployment_model.load_from_disk(model_path=model_path, device=device)
 
-        return model, "torch"
+        log.info("Imported %s model", deployment_model.__class__.__name__)
+
+        return deployment_model
 
     raise ValueError(f"Unable to load model with extension {file_extension}, valid extensions are: ['.pt', 'pth']")
+
+
+# This may be better as a dict?
+def get_export_extension(export_type: str) -> str:
+    """Get the extension of the exported model.
+
+    Args:
+        export_type: The type of the exported model.
+
+    Returns:
+        The extension of the exported model.
+    """
+    if export_type == "onnx":
+        extension = "onnx"
+    elif export_type == "torchscript":
+        extension = "pt"
+    elif export_type == "pytorch":
+        extension = "pth"
+    else:
+        raise ValueError(f"Unsupported export type {export_type}")
+
+    return extension
diff --git a/quadra/utils/logger.py b/quadra/utils/logger.py
new file mode 100644
index 00000000..46fdd53b
--- /dev/null
+++ b/quadra/utils/logger.py
@@ -0,0 +1,15 @@
+import logging
+
+from pytorch_lightning.utilities import rank_zero_only
+
+
+def get_logger(name=__name__) -> logging.Logger:
+    """Initializes multi-GPU-friendly python logger."""
+    logger = logging.getLogger(name)
+
+    # this ensures all logging levels get marked with the rank zero decorator
+    # otherwise logs would get multiplied for each GPU process in multi-GPU setup
+    for level in ("debug", "info", "warning", "error", "exception", "fatal", "critical"):
+        setattr(logger, level, rank_zero_only(getattr(logger, level)))
+
+    return logger
diff --git a/quadra/utils/models.py b/quadra/utils/models.py
index baa5dc78..d078997b 100644
--- a/quadra/utils/models.py
+++ b/quadra/utils/models.py
@@ -14,9 +14,12 @@
 from timm.models.vision_transformer import Mlp
 from torch import nn
 
+from quadra.models.evaluation import BaseEvaluationModel, TorchEvaluationModel, TorchscriptEvaluationModel
 from quadra.utils import utils
 from quadra.utils.vit_explainability import VitAttentionGradRollout
 
+log = utils.get_logger(__name__)
+
 
 def net_hat(input_size: int, output_size: int) -> torch.nn.Sequential:
     """Create a linear layer with input and output neurons.
@@ -77,7 +80,7 @@ def init_weights(m):
 
 
 def get_feature(
-    feature_extractor: torch.nn.Module,
+    feature_extractor: Union[torch.nn.Module, BaseEvaluationModel],
     dl: torch.utils.data.DataLoader,
     iteration_over_training: int = 1,
     gradcam: bool = False,
@@ -101,11 +104,19 @@ def get_feature(
             labels: input_labels
             grayscale_cams: Gradcam output maps, None if gradcam arg is False
     """
+    if isinstance(feature_extractor, (TorchEvaluationModel, TorchscriptEvaluationModel)):
+        # If we are working with torch based evaluation models we need to extract the model
+        feature_extractor = feature_extractor.model
+    else:
+        gradcam = False
+
     feature_extractor.eval()
 
     # Setup gradcam
     if gradcam:
-        if isinstance(feature_extractor.features_extractor, timm.models.resnet.ResNet):
+        if not hasattr(feature_extractor, "features_extractor"):
+            gradcam = False
+        elif isinstance(feature_extractor.features_extractor, timm.models.resnet.ResNet):
             target_layers = [feature_extractor.features_extractor.layer4[-1]]
             cam = GradCAM(
                 model=feature_extractor,
@@ -121,22 +132,27 @@ def get_feature(
                 example_input=None if input_shape is None else torch.randn(1, *input_shape),
             )
         else:
-            log = utils.get_logger(__name__)
-            log.warning("Gradcam not implemented for this backbone, it will not be computed")
             gradcam = False
 
+        if not gradcam:
+            log.warning("Gradcam not implemented for this backbone, it will not be computed")
+
     # Extract features from data
 
     for iteration in range(iteration_over_training):
         for i, b in enumerate(tqdm.tqdm(dl)):
             x1, y1 = b
-            # Move input to the correct device
-            x1 = x1.to(next(feature_extractor.parameters()).device)
+
+            if hasattr(feature_extractor, "parameters"):
+                # Move input to the correct device
+                x1 = x1.to(next(feature_extractor.parameters()).device)
+
             if gradcam:
                 y_hat = cast(
                     Union[List[torch.Tensor], Tuple[torch.Tensor], torch.Tensor], feature_extractor(x1).detach()
                 )
-                if is_vision_transformer(feature_extractor.features_extractor):  # type: ignore[arg-type]
+                # mypy can't detect that gradcam is true only if we have a features_extractor
+                if is_vision_transformer(feature_extractor.features_extractor):  # type: ignore[union-attr, arg-type]
                     grayscale_cam_low_res = grad_rollout(
                         input_tensor=x1, targets_list=y1
                     )  # TODO: We are using labels (y1) but it would be better to use preds
@@ -146,9 +162,8 @@ def get_feature(
                     grayscale_cam = ndimage.zoom(grayscale_cam_low_res, zoom_factors, order=1)
                 else:
                     grayscale_cam = cam(input_tensor=x1, targets=None)
-                feature_extractor.zero_grad(set_to_none=True)
+                feature_extractor.zero_grad(set_to_none=True)  # type: ignore[union-attr]
                 torch.cuda.empty_cache()
-
             else:
                 with torch.no_grad():
                     y_hat = cast(Union[List[torch.Tensor], Tuple[torch.Tensor], torch.Tensor], feature_extractor(x1))
diff --git a/quadra/utils/utils.py b/quadra/utils/utils.py
index d10636a9..bc226933 100644
--- a/quadra/utils/utils.py
+++ b/quadra/utils/utils.py
@@ -31,8 +31,6 @@
 
 IMAGE_EXTENSIONS: List[str] = [".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".pbm", ".pgm", ".ppm", ".pxm", ".pnm"]
 
-IMAGE_EXTENSIONS: List[str] = [".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".pbm", ".pgm", ".ppm", ".pxm", ".pnm"]
-
 
 def get_logger(name=__name__) -> logging.Logger:
     """Initializes multi-GPU-friendly python logger."""
@@ -181,11 +179,19 @@ def log_hyperparameters(
         )
         == 0
     ):
-        hparams["git/commit"] = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("ascii").strip()
-        hparams["git/branch"] = (
-            subprocess.check_output(["git", "rev-parse", "--abbrev-ref", "HEAD"]).decode("ascii").strip()
-        )
-        hparams["git/remote"] = subprocess.check_output(["git", "remote", "get-url", "origin"]).decode("ascii").strip()
+        try:
+            hparams["git/commit"] = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("ascii").strip()
+            hparams["git/branch"] = (
+                subprocess.check_output(["git", "rev-parse", "--abbrev-ref", "HEAD"]).decode("ascii").strip()
+            )
+            hparams["git/remote"] = (
+                subprocess.check_output(["git", "remote", "get-url", "origin"]).decode("ascii").strip()
+            )
+        except subprocess.CalledProcessError:
+            log.warning(
+                "Could not get git commit, branch or remote information, the repository might not have any commits yet "
+                "or it might be initialized wrongly."
+            )
     else:
         log.warning("Could not find git repository, skipping git commit and branch info")
 
@@ -270,7 +276,9 @@ def finish(
             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")
+                        model = quadra_export.import_deployment_model(
+                            model_path, device="cpu", inference_config=config.inference
+                        ).model
 
                         input_size = model_json["input_size"]
 


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index deb7f18b..22ed32de 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -37,9 +37,8 @@ jobs:
       - name: Install Package
         run: |
           python -m pip install -U pip
-          python -m pip install -e ".[test]" --no-cache-dir
+          python -m pip install -e ".[test,onnx]" --no-cache-dir
 
       - name: Run Tests
         run: |
           python -m pytest -v --disable-pytest-warnings --strict-markers --color=yes -m "not slow"
- 
diff --git a/quadra/utils/tests/fixtures/dataset/__init__.py b/quadra/utils/tests/fixtures/dataset/__init__.py
index a8a47e51..49a53bd0 100644
--- a/quadra/utils/tests/fixtures/dataset/__init__.py
+++ b/quadra/utils/tests/fixtures/dataset/__init__.py
@@ -10,6 +10,7 @@
     classification_patch_dataset,
     multilabel_classification_dataset,
 )
+from .imagenette import imagenette_dataset
 from .segmentation import (
     SegmentationDatasetArguments,
     base_binary_segmentation_dataset,
@@ -29,6 +30,7 @@
     "multilabel_classification_dataset",
     "ClassificationMultilabelDatasetArguments",
     "base_anomaly_dataset",
+    "imagenette_dataset",
     "base_classification_dataset",
     "base_patch_classification_dataset",
     "base_binary_segmentation_dataset",
diff --git a/quadra/utils/tests/fixtures/dataset/classification.py b/quadra/utils/tests/fixtures/dataset/classification.py
index e0a478fc..307e3b3c 100644
--- a/quadra/utils/tests/fixtures/dataset/classification.py
+++ b/quadra/utils/tests/fixtures/dataset/classification.py
@@ -374,7 +374,7 @@ def classification_patch_dataset(
     params=[
         ClassificationPatchDatasetArguments(
             **{
-                "samples": [15, 20, 25],
+                "samples": [5, 5, 5],
                 "classes": ["bg", "a", "b"],
                 "patch_size": [16, 20],
                 "overlap": 0.6,
@@ -389,7 +389,7 @@ def base_patch_classification_dataset(
 ) -> Tuple[str, ClassificationDatasetArguments, Dict[str, int]]:
     """Generate a classification patch dataset with the following parameters:
         - 3 classes named bg, a and b
-        - 15, 20 and 25 samples for each class
+        - 5, 5 and 5 samples for each class
         - patch size of 16x20
         - 60% overlap
         - 10% validation set
diff --git a/quadra/utils/tests/fixtures/dataset/imagenette.py b/quadra/utils/tests/fixtures/dataset/imagenette.py
new file mode 100644
index 00000000..86683aed
--- /dev/null
+++ b/quadra/utils/tests/fixtures/dataset/imagenette.py
@@ -0,0 +1,53 @@
+import shutil
+from pathlib import Path
+
+import cv2
+import pytest
+
+from quadra.utils.tests.helpers import _random_image
+
+
+def _build_imagenette_dataset(tmp_path: Path, classes: int, class_samples: int) -> str:
+    """Generate imagenette dataset in the format required by efficient_ad model.
+
+    Args:
+        tmp_path: Path to temporary directory
+        classes: Number of mock imagenette classes
+        class_samples: Number of samples for each mock imagenette class
+
+    Returns:
+        Path to imagenette dataset
+    """
+    parent_path = tmp_path / "imagenette2"
+    parent_path.mkdir()
+    train_path = parent_path / "train"
+    train_path.mkdir()
+    val_path = parent_path / "val"
+    val_path.mkdir()
+
+    for split in [train_path, val_path]:
+        for i in range(classes):
+            cl_path = split / f"class_{i}"
+            cl_path.mkdir()
+            for j in range(class_samples):
+                image = _random_image()
+                image_path = cl_path / f"fake_{j}.png"
+                cv2.imwrite(str(image_path), image)
+
+    return parent_path
+
+
+@pytest.fixture
+def imagenette_dataset(tmp_path: Path) -> str:
+    """Generate a mock imagenette dataset to test efficient_ad model
+
+    Args:
+        tmp_path: Path to temporary directory
+        request: Pytest SubRequest object
+    Yields:
+        Path to imagenette dataset folder
+    """
+    yield _build_imagenette_dataset(tmp_path, classes=3, class_samples=3)
+
+    if tmp_path.exists():
+        shutil.rmtree(tmp_path)
diff --git a/quadra/utils/tests/fixtures/models/__init__.py b/quadra/utils/tests/fixtures/models/__init__.py
new file mode 100644
index 00000000..38502b2e
--- /dev/null
+++ b/quadra/utils/tests/fixtures/models/__init__.py
@@ -0,0 +1,3 @@
+from .anomaly import *
+from .classification import *
+from .segmentation import *
diff --git a/quadra/utils/tests/fixtures/models/anomaly.py b/quadra/utils/tests/fixtures/models/anomaly.py
new file mode 100644
index 00000000..8281ea3a
--- /dev/null
+++ b/quadra/utils/tests/fixtures/models/anomaly.py
@@ -0,0 +1,69 @@
+import pytest
+import torch
+from anomalib.models.draem.torch_model import DraemModel
+from anomalib.models.padim.torch_model import PadimModel
+from anomalib.models.patchcore.torch_model import PatchcoreModel
+
+
+@pytest.fixture
+def padim_resnet18():
+    """Yield a padim model with resnet18 encoder."""
+    yield PadimModel(
+        input_size=[224, 224],  # TODO: This is hardcoded may be not a good idea
+        backbone="resnet18",
+        layers=["layer1", "layer2", "layer3"],
+        pretrained_weights=None,
+        tied_covariance=False,
+        pre_trained=False,
+    )
+
+
+@torch.inference_mode()
+def _initialize_patchcore_model(patchcore_model: PatchcoreModel, coreset_sampling_ratio: float = 0.1) -> PatchcoreModel:
+    """Initialize a Patchcore model by simulating a training step.
+
+    Args:
+        patchcore_model: Patchcore model to initialize
+        coreset_sampling_ratio: Coreset sampling ratio to use for the initialization
+
+    Returns:
+        Patchcore model with initialized memory bank
+    """
+    with torch.no_grad():
+        training_features = None
+        random_input = torch.rand([1, 3, *patchcore_model.input_size])
+
+        if training_features is None:
+            training_features = patchcore_model(random_input)
+        else:
+            training_features = torch.cat([training_features, patchcore_model(random_input)], dim=0)
+
+        patchcore_model.eval()
+        patchcore_model.subsample_embedding(training_features, sampling_ratio=coreset_sampling_ratio)
+
+        # Simulate a memory bank with 5 images, at the current stage patchcore onnx export is not handling
+        # large memory banks well, so we are using a small one for the benchmark
+        memory_bank_number, memory_bank_n_features = patchcore_model.memory_bank.shape
+        patchcore_model.memory_bank = torch.rand([5 * memory_bank_number, memory_bank_n_features])
+        patchcore_model.train()
+
+    return patchcore_model
+
+
+@pytest.fixture
+def patchcore_resnet18():
+    """Yield a patchcore model with resnet18 encoder."""
+    model = PatchcoreModel(
+        input_size=[224, 224],  # TODO: This is hardcoded may be not a good idea
+        backbone="resnet18",
+        layers=["layer2", "layer3"],
+        pre_trained=False,
+    )
+
+    yield _initialize_patchcore_model(model)
+
+
+@pytest.fixture
+def draem():
+    """Yield a draem model."""
+    yield DraemModel()
diff --git a/quadra/utils/tests/fixtures/models/classification.py b/quadra/utils/tests/fixtures/models/classification.py
new file mode 100644
index 00000000..94b841da
--- /dev/null
+++ b/quadra/utils/tests/fixtures/models/classification.py
@@ -0,0 +1,45 @@
+import pytest
+
+from quadra.models.classification import TimmNetworkBuilder, TorchHubNetworkBuilder
+
+
+@pytest.fixture
+def resnet18():
+    """Yield a resnet18 model."""
+    yield TimmNetworkBuilder("resnet18", pretrained=False, freeze=True, exportable=True)
+
+
+@pytest.fixture
+def resnet50():
+    """Yield a resnet50 model."""
+    yield TimmNetworkBuilder("resnet50", pretrained=False, freeze=True, exportable=True)
+
+
+@pytest.fixture
+def vit_tiny_patch16_224():
+    """Yield a vit_tiny_patch16_224 model."""
+    yield TimmNetworkBuilder("vit_tiny_patch16_224", pretrained=False, freeze=True, exportable=True)
+
+
+@pytest.fixture
+def dino_vits8():
+    """Yield a dino_vits8 model."""
+    yield TorchHubNetworkBuilder(
+        repo_or_dir="facebookresearch/dino:main",
+        model_name="dino_vits8",
+        pretrained=False,
+        freeze=True,
+        exportable=True,
+    )
+
+
+@pytest.fixture
+def dino_vitb8():
+    """Yield a dino_vitb8 model."""
+    yield TorchHubNetworkBuilder(
+        repo_or_dir="facebookresearch/dino:main",
+        model_name="dino_vitb8",
+        pretrained=False,
+        freeze=True,
+        exportable=True,
+    )
diff --git a/quadra/utils/tests/fixtures/models/segmentation.py b/quadra/utils/tests/fixtures/models/segmentation.py
new file mode 100644
index 00000000..4b22eb37
--- /dev/null
+++ b/quadra/utils/tests/fixtures/models/segmentation.py
@@ -0,0 +1,33 @@
+import pytest
+
+from quadra.modules.backbone import create_smp_backbone
+
+
+@pytest.fixture
+def smp_resnet18_unet():
+    """Yield a unet with resnet18 encoder."""
+    yield create_smp_backbone(
+        arch="unet",
+        encoder_name="resnet18",
+        encoder_weights=None,
+        encoder_depth=5,
+        freeze_encoder=True,
+        in_channels=3,
+        num_classes=1,
+        activation=None,
+    )
+
+
+@pytest.fixture
+def smp_resnet18_unetplusplus():
+    """Yield a unetplusplus with resnet18 encoder."""
+    yield create_smp_backbone(
+        arch="unetplusplus",
+        encoder_name="resnet18",
+        encoder_weights=None,
+        encoder_depth=5,
+        freeze_encoder=True,
+        in_channels=3,
+        num_classes=1,
+        activation=None,
+    )
diff --git a/quadra/utils/tests/helpers.py b/quadra/utils/tests/helpers.py
index 45b68412..40bb2cc8 100644
--- a/quadra/utils/tests/helpers.py
+++ b/quadra/utils/tests/helpers.py
@@ -3,10 +3,12 @@
 from typing import List, Tuple
 
 import numpy as np
+import torch
 from hydra import compose, initialize_config_module
 from hydra.core.hydra_config import HydraConfig
 
 from quadra.main import main
+from quadra.utils.export import get_export_extension
 
 
 # taken from hydra unit tests
@@ -28,3 +30,40 @@ def execute_quadra_experiment(overrides: List[str], experiment_path: Path) -> No
         HydraConfig.instance().set_config(cfg)
 
         main(cfg)
+
+
+def check_deployment_model(export_type: str):
+    """Check that the runtime model is present and valid.
+
+    Args:
+        export_type: The type of the exported model.
+    """
+    extension = get_export_extension(export_type)
+
+    assert os.path.exists(f"deployment_model/model.{extension}")
+    assert os.path.exists("deployment_model/model.json")
+
+
+def get_quadra_test_device():
+    """Get the device to use for the tests. If the QUADRA_TEST_DEVICE environment variable is set, it is used."""
+    return os.environ.get("QUADRA_TEST_DEVICE", "cpu")
+
+
+def setup_trainer_for_lightning() -> List[str]:
+    """Setup trainer for lightning depending on the device. If cuda is used, the device index is also set.
+    If cpu is used, the trainer is set to lightning_cpu.
+
+    Returns:
+        A list of overrides for the trainer.
+    """
+    overrides = []
+    device = get_quadra_test_device()
+    torch_device = torch.device(device)
+    if torch_device.type == "cuda":
+        device_index = torch_device.index
+        overrides.append("trainer=lightning_gpu")
+        overrides.append(f"trainer.devices=[{device_index}]")
+    else:
+        overrides.append("trainer=lightning_cpu")
+
+    return overrides
diff --git a/tests/conftest.py b/tests/conftest.py
index 67a3277a..c668bd3d 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -2,6 +2,11 @@
 from pathlib import Path
 
 import pytest
+import torch
+
+
+def pytest_addoption(parser):
+    parser.addoption("--device", action="store", default="cpu", help="device to run tests on")
 
 
 @pytest.fixture(autouse=True)
@@ -13,7 +18,15 @@ def change_test_dir(tmp_path: Path, monkeypatch):
     monkeypatch.chdir(test_working_directory)
 
 
+@pytest.fixture(scope="session")
+def device(pytestconfig):
+    return pytestconfig.getoption("device")
+
+
 @pytest.fixture(autouse=True)
-def hyde_gpu():
-    """Hyde GPU for tests."""
-    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+def setup_devices(device: str):
+    """Set the device to run tests on."""
+    torch_device = torch.device(device)
+    os.environ["QUADRA_TEST_DEVICE"] = device
+    if torch_device.type != "cuda":
+        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
diff --git a/tests/models/test_export.py b/tests/models/test_export.py
new file mode 100644
index 00000000..8cfb6e4e
--- /dev/null
+++ b/tests/models/test_export.py
@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Sequence
+
+import pytest
+import torch
+from omegaconf import DictConfig
+from torch import nn
+
+from quadra.utils.export import export_onnx_model, export_torchscript_model, import_deployment_model
+from quadra.utils.tests.fixtures.models import (  # noqa
+    dino_vitb8,
+    dino_vits8,
+    draem,
+    padim_resnet18,
+    patchcore_resnet18,
+    resnet18,
+    resnet50,
+    smp_resnet18_unet,
+    smp_resnet18_unetplusplus,
+    vit_tiny_patch16_224,
+)
+
+try:
+    import onnx  # noqa
+    import onnxruntime  # noqa
+    import onnxsim  # noqa
+
+    ONNX_AVAILABLE = True
+except ImportError:
+    ONNX_AVAILABLE = False
+
+ONNX_CONFIG = DictConfig(
+    {
+        "opset_version": 16,
+        "do_constant_folding": True,
+        "export_params": True,
+        "simplify": False,
+    }
+)
+
+
+@torch.inference_mode()
+def check_export_model_outputs(tmp_path: Path, model: nn.Module, export_types: list[str], input_shapes: tuple[Any]):
+    exported_models = {}
+
+    for export_type in export_types:
+        if export_type == "torchscript":
+            out = export_torchscript_model(
+                model=model,
+                input_shapes=input_shapes,
+                output_path=tmp_path,
+                half_precision=False,
+            )
+
+            torchscript_model_path, input_shapes = out
+            exported_models[export_type] = torchscript_model_path
+        else:
+            out = export_onnx_model(
+                model=model,
+                output_path=tmp_path,
+                onnx_config=ONNX_CONFIG,
+                input_shapes=input_shapes,
+                half_precision=False,
+            )
+
+            onnx_model_path, input_shapes = out
+            exported_models[export_type] = onnx_model_path
+
+    inference_config = DictConfig({"onnx": {}, "torchscript": {}})
+
+    models = []
+    for export_type, model_path in exported_models.items():
+        model = import_deployment_model(model_path=model_path, inference_config=inference_config, device="cpu")
+        models.append(model)
+
+    inp = torch.rand((1, *input_shapes[0]), dtype=torch.float32)
+
+    outputs = []
+
+    for model in models:
+        outputs.append(model(inp))
+
+    for i in range(len(outputs) - 1):
+        if isinstance(outputs[i], Sequence):
+            for j in range(len(outputs[i])):
+                assert torch.allclose(outputs[i][j], outputs[i + 1][j], atol=1e-5)
+        else:
+            assert torch.allclose(outputs[i], outputs[i + 1], atol=1e-5)
+
+
+@pytest.mark.skipif(not ONNX_AVAILABLE, reason="ONNX not available")
+@pytest.mark.parametrize(
+    "model",
+    [
+        pytest.lazy_fixture("dino_vitb8"),
+        pytest.lazy_fixture("dino_vits8"),
+        pytest.lazy_fixture("resnet18"),
+        pytest.lazy_fixture("resnet50"),
+        pytest.lazy_fixture("vit_tiny_patch16_224"),
+    ],
+)
+def test_classification_models_export(tmp_path: Path, model: nn.Module):
+    export_types = ["onnx", "torchscript"]
+
+    input_shapes = [(3, 224, 224)]
+
+    check_export_model_outputs(tmp_path=tmp_path, model=model, export_types=export_types, input_shapes=input_shapes)
+
+
+@pytest.mark.skipif(not ONNX_AVAILABLE, reason="ONNX not available")
+@pytest.mark.parametrize(
+    "model",
+    [
+        pytest.lazy_fixture("smp_resnet18_unet"),
+        pytest.lazy_fixture("smp_resnet18_unetplusplus"),
+    ],
+)
+def test_segmentation_models_export(tmp_path: Path, model: nn.Module):
+    export_types = ["onnx", "torchscript"]
+
+    input_shapes = [(3, 224, 224)]
+
+    check_export_model_outputs(tmp_path=tmp_path, model=model, export_types=export_types, input_shapes=input_shapes)
+
+
+@pytest.mark.skipif(not ONNX_AVAILABLE, reason="ONNX not available")
+@pytest.mark.parametrize(
+    "model",
+    [
+        pytest.lazy_fixture("padim_resnet18"),
+        pytest.lazy_fixture("patchcore_resnet18"),
+        pytest.lazy_fixture("draem"),
+    ],
+)
+def test_anomaly_detection_models_export(tmp_path: Path, model: nn.Module):
+    export_types = ["onnx", "torchscript"]
+
+    input_shapes = [(3, 224, 224)]
+
+    check_export_model_outputs(tmp_path=tmp_path, model=model, export_types=export_types, input_shapes=input_shapes)
diff --git a/tests/tasks/test_anomaly.py b/tests/tasks/test_anomaly.py
index 27a44549..dec4ccb9 100644
--- a/tests/tasks/test_anomaly.py
+++ b/tests/tasks/test_anomaly.py
@@ -2,14 +2,24 @@
 import os
 import shutil
 from pathlib import Path
+from typing import List
 
 import pytest
 
-from quadra.utils.tests.fixtures import base_anomaly_dataset
-from quadra.utils.tests.helpers import execute_quadra_experiment
+from quadra.utils.export import get_export_extension
+from quadra.utils.tests.fixtures import base_anomaly_dataset, imagenette_dataset
+from quadra.utils.tests.helpers import check_deployment_model, execute_quadra_experiment, setup_trainer_for_lightning
+
+try:
+    import onnx  # noqa
+    import onnxruntime  # noqa
+    import onnxsim  # noqa
+
+    ONNX_AVAILABLE = True
+except ImportError:
+    ONNX_AVAILABLE = False
 
 BASE_EXPERIMENT_OVERRIDES = [
-    "trainer=lightning_cpu",
     "trainer.devices=1",
     "datamodule.num_workers=1",
     "datamodule.train_batch_size=1",
@@ -20,19 +30,7 @@
     "~logger.mlflow",
 ]
 
-
-def _check_deployment_model(invert: bool = False):
-    """Check that the runtime model is present and valid.
-
-    Args:
-        invert: If true check that the runtime model is not present.
-    """
-    if invert:
-        assert not os.path.exists("deployment_model/model.pt")
-        assert not os.path.exists("deployment_model/model.json")
-    else:
-        assert os.path.exists("deployment_model/model.pt")
-        assert os.path.exists("deployment_model/model.json")
+BASE_EXPORT_TYPES = ["torchscript"] if not ONNX_AVAILABLE else ["torchscript", "onnx"]
 
 
 def _check_report(invert: bool = False):
@@ -51,7 +49,10 @@ def _check_report(invert: bool = False):
         assert os.path.exists("avg_score_by_label.csv")
 
 
-def _run_inference_experiment(data_path: str, train_path: str, test_path: str):
+def _run_inference_experiment(data_path: str, train_path: str, test_path: str, export_type: str):
+    """Run an inference experiment for the given export type."""
+    extension = get_export_extension(export_type)
+
     test_overrides = [
         "task.device=cpu",
         "experiment=base/anomaly/inference",
@@ -59,14 +60,29 @@ def _run_inference_experiment(data_path: str, train_path: str, test_path: str):
         "datamodule.num_workers=1",
         "datamodule.test_batch_size=32",
         "logger=csv",
-        f"task.model_path={os.path.join(train_path, 'deployment_model', 'model.pt')}",
+        f"task.model_path={os.path.join(train_path, 'deployment_model', f'model.{extension}')}",
     ]
 
     execute_quadra_experiment(overrides=test_overrides, experiment_path=test_path)
 
 
+def run_inference_experiments(data_path: str, train_path: str, test_path: str, export_types: List[str]):
+    """Run inference experiments for the given export types."""
+    for export_type in export_types:
+        cwd = os.getcwd()
+        check_deployment_model(export_type=export_type)
+
+        _run_inference_experiment(
+            data_path=data_path, train_path=train_path, test_path=test_path, export_type=export_type
+        )
+
+        # Change back to the original working directory
+        os.chdir(cwd)
+
+
 @pytest.mark.parametrize("task", ["classification", "segmentation"])
-def test_padim_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_dataset, task: str):
+def test_padim(tmp_path: Path, base_anomaly_dataset: base_anomaly_dataset, task: str):
+    """Test the training and evaluation of the PADIM model."""
     data_path, _ = base_anomaly_dataset
 
     train_path = tmp_path / "train"
@@ -77,23 +93,28 @@ 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}",
+        f"export.types=[{','.join(BASE_EXPORT_TYPES)}]",
     ]
+    trainer_overrides = setup_trainer_for_lightning()
     overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=train_path)
 
     assert os.path.exists("checkpoints/final_model.ckpt")
 
-    _check_deployment_model()
     _check_report()
 
-    _run_inference_experiment(data_path, train_path, test_path)
+    run_inference_experiments(
+        data_path=data_path, train_path=train_path, test_path=test_path, export_types=BASE_EXPORT_TYPES
+    )
 
     shutil.rmtree(tmp_path)
 
 
 @pytest.mark.parametrize("task", ["classification", "segmentation"])
-def test_patchcore_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_dataset, task: str):
+def test_patchcore(tmp_path: Path, base_anomaly_dataset: base_anomaly_dataset, task: str):
+    """Test the training and evaluation of the PatchCore model."""
     data_path, _ = base_anomaly_dataset
 
     train_path = tmp_path / "train"
@@ -104,24 +125,70 @@ 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}",
+        f"export.types=[{','.join(BASE_EXPORT_TYPES)}]",
+    ]
+    trainer_overrides = setup_trainer_for_lightning()
+    overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
+
+    execute_quadra_experiment(overrides=overrides, experiment_path=train_path)
+
+    assert os.path.exists("checkpoints/final_model.ckpt")
+
+    _check_report()
+
+    run_inference_experiments(
+        data_path=data_path, train_path=train_path, test_path=test_path, export_types=BASE_EXPORT_TYPES
+    )
+
+    shutil.rmtree(tmp_path)
+
+
+@pytest.mark.parametrize("task", ["classification", "segmentation"])
+def test_efficientad(
+    tmp_path: Path, base_anomaly_dataset: base_anomaly_dataset, imagenette_dataset: imagenette_dataset, task: str
+):
+    """Test the training and evaluation of the EfficientAD model."""
+    data_path, _ = base_anomaly_dataset
+    imagenette_path = imagenette_dataset
+
+    train_path = tmp_path / "train"
+    test_path = tmp_path / "test"
+
+    overrides = [
+        "experiment=base/anomaly/efficient_ad",
+        f"datamodule.data_path={data_path}",
+        "transforms.input_height=256",
+        "transforms.input_width=256",
+        "model.model.train_batch_size=1",
+        "datamodule.test_batch_size=1",
+        "model.model.image_size=[256, 256]",
+        "trainer.check_val_every_n_epoch= ${trainer.max_epochs}",
+        f"model.model.imagenette_dir= {imagenette_path}",
+        f"model.dataset.task={task}",
+        f"export.types=[{','.join(BASE_EXPORT_TYPES)}]",
     ]
+    trainer_overrides = setup_trainer_for_lightning()
     overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=train_path)
 
     assert os.path.exists("checkpoints/final_model.ckpt")
 
-    _check_deployment_model()
     _check_report()
 
-    _run_inference_experiment(data_path, train_path, test_path)
+    run_inference_experiments(
+        data_path=data_path, train_path=train_path, test_path=test_path, export_types=BASE_EXPORT_TYPES
+    )
 
     shutil.rmtree(tmp_path)
 
 
 @pytest.mark.slow
 @pytest.mark.parametrize("task", ["classification", "segmentation"])
-def test_cflow_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_dataset, task: str):
+def test_cflow(tmp_path: Path, base_anomaly_dataset: base_anomaly_dataset, task: str):
+    """Test the training and evaluation of the cflow model."""
     data_path, _ = base_anomaly_dataset
 
     overrides = [
@@ -129,24 +196,26 @@ 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_config=null",
+        "export.types=[]",
     ]
+    trainer_overrides = setup_trainer_for_lightning()
     overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=tmp_path)
 
     assert os.path.exists("checkpoints/final_model.ckpt")
 
-    _check_deployment_model(invert=True)  # cflow does not support runtime model
     _check_report()
 
-    # cflow does not support inference with jitted model
+    # cflow does not support exporting to torchscript and onnx at the moment
     shutil.rmtree(tmp_path)
 
 
 @pytest.mark.slow
 @pytest.mark.parametrize("task", ["classification", "segmentation"])
-def test_csflow_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_dataset, task: str):
+def test_csflow(tmp_path: Path, base_anomaly_dataset: base_anomaly_dataset, task: str):
+    """Test the training and evaluation of the csflow model."""
     data_path, _ = base_anomaly_dataset
     train_path = tmp_path / "train"
     test_path = tmp_path / "test"
@@ -155,51 +224,63 @@ 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}",
+        f"export.types=[{','.join(BASE_EXPORT_TYPES)}]",
     ]
+    trainer_overrides = setup_trainer_for_lightning()
     overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=train_path)
 
     assert os.path.exists("checkpoints/final_model.ckpt")
 
-    _check_deployment_model()
     _check_report()
 
-    _run_inference_experiment(data_path, train_path, test_path)
+    run_inference_experiments(
+        data_path=data_path, train_path=train_path, test_path=test_path, export_types=BASE_EXPORT_TYPES
+    )
 
     shutil.rmtree(tmp_path)
 
 
 @pytest.mark.slow
 @pytest.mark.parametrize("task", ["classification", "segmentation"])
-def test_fastflow_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_dataset, task: str):
+def test_fastflow(tmp_path: Path, base_anomaly_dataset: base_anomaly_dataset, task: str):
+    """Test the training and evaluation of the fastflow model."""
     data_path, _ = base_anomaly_dataset
     train_path = tmp_path / "train"
     test_path = tmp_path / "test"
 
+    export_types = ["torchscript"]  # fastflow does not support exporting to onnx at the moment
+
     overrides = [
         "experiment=base/anomaly/fastflow",
         f"datamodule.data_path={data_path}",
         "model.model.backbone=resnet18",
         f"model.dataset.task={task}",
+        f"export.types=[{','.join(export_types)}]",
     ]
+    trainer_overrides = setup_trainer_for_lightning()
     overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=train_path)
 
     assert os.path.exists("checkpoints/final_model.ckpt")
 
-    _check_deployment_model()
     _check_report()
 
-    _run_inference_experiment(data_path, train_path, test_path)
+    run_inference_experiments(
+        data_path=data_path, train_path=train_path, test_path=test_path, export_types=export_types
+    )
 
     shutil.rmtree(tmp_path)
 
 
 @pytest.mark.slow
 @pytest.mark.parametrize("task", ["classification", "segmentation"])
-def test_draem_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_dataset, task: str):
+def test_draem(tmp_path: Path, base_anomaly_dataset: base_anomaly_dataset, task: str):
+    """Test the training and evaluation of the draem model."""
     data_path, _ = base_anomaly_dataset
     train_path = tmp_path / "train"
     test_path = tmp_path / "test"
@@ -208,16 +289,20 @@ 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}",
+        f"export.types=[{','.join(BASE_EXPORT_TYPES)}]",
     ]
+    trainer_overrides = setup_trainer_for_lightning()
     overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=train_path)
 
     assert os.path.exists("checkpoints/final_model.ckpt")
 
-    _check_deployment_model()
     _check_report()
 
-    _run_inference_experiment(data_path, train_path, test_path)
+    run_inference_experiments(
+        data_path=data_path, train_path=train_path, test_path=test_path, export_types=BASE_EXPORT_TYPES
+    )
 
     shutil.rmtree(tmp_path)
diff --git a/tests/tasks/test_classification.py b/tests/tasks/test_classification.py
index 712e1e6f..398c82af 100644
--- a/tests/tasks/test_classification.py
+++ b/tests/tasks/test_classification.py
@@ -2,45 +2,84 @@
 import os
 import shutil
 from pathlib import Path
+from typing import List
 
 import pytest
 
+from quadra.utils.export import get_export_extension
 from quadra.utils.tests.fixtures import (
     base_classification_dataset,
     base_multilabel_classification_dataset,
     base_patch_classification_dataset,
 )
-from quadra.utils.tests.helpers import execute_quadra_experiment
+from quadra.utils.tests.helpers import (
+    check_deployment_model,
+    execute_quadra_experiment,
+    get_quadra_test_device,
+    setup_trainer_for_lightning,
+)
+
+try:
+    import onnx  # noqa
+    import onnxruntime  # noqa
+    import onnxsim  # noqa
+
+    ONNX_AVAILABLE = True
+except ImportError:
+    ONNX_AVAILABLE = False
 
 BASE_EXPERIMENT_OVERRIDES = [
     "datamodule.num_workers=1",
     "logger=csv",
 ]
 
+BASE_EXPORT_TYPES = ["pytorch", "torchscript"] if not ONNX_AVAILABLE else ["pytorch", "torchscript", "onnx"]
 
-def test_train_sklearn_classification(tmp_path: Path, base_classification_dataset: base_classification_dataset):
-    data_path, _ = base_classification_dataset
 
-    overrides = [
-        "experiment=base/classification/sklearn_classification",
-        f"datamodule.data_path={data_path}",
-        "backbone=resnet18",
-        "task.device=cpu",
-    ] + BASE_EXPERIMENT_OVERRIDES
+def _run_inference_experiment(
+    test_overrides: List[str], data_path: str, train_path: str, test_path: str, export_type: str
+):
+    """Run an inference experiment for the given export type."""
+    extension = get_export_extension(export_type)
 
-    execute_quadra_experiment(overrides=overrides, experiment_path=tmp_path)
+    test_overrides.append(f"datamodule.data_path={data_path}")
+    test_overrides.append(f"task.model_path={os.path.join(train_path, 'deployment_model', f'model.{extension}')}")
 
-    shutil.rmtree(tmp_path)
+    execute_quadra_experiment(overrides=test_overrides, experiment_path=test_path)
+
+
+def run_inference_experiments(
+    test_overrides: List[str], data_path: str, train_path: str, test_path: str, export_types: List[str]
+):
+    """Run inference experiments for the given export types."""
+    for export_type in export_types:
+        cwd = os.getcwd()
+        check_deployment_model(export_type=export_type)
+
+        _run_inference_experiment(
+            test_overrides=test_overrides,
+            data_path=data_path,
+            train_path=train_path,
+            test_path=test_path,
+            export_type=export_type,
+        )
 
+        # Change back to the original working directory
+        os.chdir(cwd)
 
-def test_inference_sklearn_classification(tmp_path: Path, base_classification_dataset: base_classification_dataset):
+
+def test_sklearn_classification(tmp_path: Path, base_classification_dataset: base_classification_dataset):
+    """Test the training and evaluation of a sklearn classification model."""
     data_path, _ = base_classification_dataset
 
+    device = get_quadra_test_device()
+
     train_overrides = [
         "experiment=base/classification/sklearn_classification",
         f"datamodule.data_path={data_path}",
         "backbone=resnet18",
-        "task.device=cpu",
+        f"task.device={device}",
+        f"export.types=[{','.join(BASE_EXPORT_TYPES)}]",
     ] + BASE_EXPERIMENT_OVERRIDES
 
     train_path = tmp_path / "train"
@@ -50,43 +89,30 @@ def test_inference_sklearn_classification(tmp_path: Path, base_classification_da
 
     execute_quadra_experiment(overrides=train_overrides, experiment_path=train_path)
 
-    trained_model_path = os.path.join(train_path, "deployment_model/model.pt")
     inference_overrides = [
         "experiment=base/classification/sklearn_classification_test",
-        f"datamodule.data_path={data_path}",
-        f"task.model_path={trained_model_path}",
         "backbone=resnet18",
-        "task.device=cpu",
+        f"task.device={device}",
         "task.gradcam=true",
     ] + BASE_EXPERIMENT_OVERRIDES
-    execute_quadra_experiment(overrides=inference_overrides, experiment_path=test_path)
-
-    shutil.rmtree(tmp_path)
-
-
-def test_train_patches(tmp_path: Path, base_patch_classification_dataset):
-    data_path, _, class_to_idx = base_patch_classification_dataset
 
-    class_to_idx_parameter = str(class_to_idx).replace(
-        "'", ""
-    )  # Remove single quotes so that it can be parsed by hydra
-
-    overrides = [
-        "experiment=base/classification/sklearn_classification_patch",
-        f"datamodule.data_path={data_path}",
-        f"datamodule.class_to_idx={class_to_idx_parameter}",
-        "trainer.iteration_over_training=1",
-        "backbone=resnet18",
-        "task.device=cpu",
-    ] + BASE_EXPERIMENT_OVERRIDES
-
-    execute_quadra_experiment(overrides=overrides, experiment_path=tmp_path)
+    run_inference_experiments(
+        test_overrides=inference_overrides,
+        data_path=data_path,
+        train_path=train_path,
+        test_path=test_path,
+        export_types=BASE_EXPORT_TYPES,
+    )
 
     shutil.rmtree(tmp_path)
 
 
-def test_inference_patches(tmp_path: Path, base_patch_classification_dataset: base_patch_classification_dataset):
+def test_sklearn_classification_patch(
+    tmp_path: Path, base_patch_classification_dataset: base_patch_classification_dataset
+):
+    """Test the training and evaluation of a sklearn classification model with patches."""
     data_path, _, class_to_idx = base_patch_classification_dataset
+    device = get_quadra_test_device()
 
     class_to_idx_parameter = str(class_to_idx).replace(
         "'", ""
@@ -105,42 +131,33 @@ def test_inference_patches(tmp_path: Path, base_patch_classification_dataset: ba
         f"datamodule.class_to_idx={class_to_idx_parameter}",
         "trainer.iteration_over_training=1",
         f"backbone={backbone}",
-        "task.device=cpu",
+        f"task.device={device}",
+        f"export.types=[{','.join(BASE_EXPORT_TYPES)}]",
     ] + BASE_EXPERIMENT_OVERRIDES
     execute_quadra_experiment(overrides=train_overrides, experiment_path=train_experiment_path)
 
-    trained_model_path = os.path.join(train_experiment_path, "deployment_model/model.pt")
     test_overrides = [
         "experiment=base/classification/sklearn_classification_patch_test",
-        f"datamodule.data_path={data_path}",
-        f"task.model_path={trained_model_path}",
         f"backbone={backbone}",
         "task.device=cpu",
     ] + BASE_EXPERIMENT_OVERRIDES
-    execute_quadra_experiment(overrides=test_overrides, experiment_path=test_experiment_path)
-
-    shutil.rmtree(tmp_path)
 
+    run_inference_experiments(
+        test_overrides=test_overrides,
+        data_path=data_path,
+        train_path=train_experiment_path,
+        test_path=test_experiment_path,
+        export_types=BASE_EXPORT_TYPES,
+    )
 
-def _run_inference_experiment(data_path: str, train_path: str, test_path: str):
-    test_overrides = [
-        "experiment=base/classification/classification_evaluation",
-        f"datamodule.data_path={data_path}",
-        "datamodule.num_workers=1",
-        "datamodule.batch_size=16",
-        "logger=csv",
-        "task.device=cpu",
-        f"task.model_path={os.path.join(train_path, 'deployment_model', 'model.pth')}",
-    ]
-
-    execute_quadra_experiment(overrides=test_overrides, experiment_path=test_path)
+    shutil.rmtree(tmp_path)
 
 
 @pytest.mark.parametrize(
     "run_test, backbone, gradcam, freeze",
     [(True, "resnet18", True, False), (False, "resnet18", False, False), (True, "dino_vits8", False, True)],
 )
-def test_train_classification(
+def test_classification(
     tmp_path: Path,
     base_classification_dataset: base_classification_dataset,
     run_test: bool,
@@ -148,15 +165,16 @@ def test_train_classification(
     gradcam: bool,
     freeze: bool,
 ):
+    """Test the training and evaluation of a torch based classification model."""
     data_path, arguments = base_classification_dataset
 
     train_path = tmp_path / "train"
     test_path = tmp_path / "test"
 
     num_classes = len(arguments.samples)
+
     overrides = [
         "experiment=base/classification/classification",
-        "trainer=lightning_cpu",
         "trainer.devices=1",
         f"datamodule.data_path={data_path}",
         f"model.num_classes={num_classes}",
@@ -166,30 +184,52 @@ def test_train_classification(
         "trainer.max_epochs=1",
         "task.report=True",
         f"task.run_test={run_test}",
-    ] + BASE_EXPERIMENT_OVERRIDES
+        f"export.types=[{','.join(BASE_EXPORT_TYPES)}]",
+    ]
+    trainer_overrides = setup_trainer_for_lightning()
+    overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=train_path)
 
-    _run_inference_experiment(data_path, train_path, test_path)
+    test_overrides = [
+        "experiment=base/classification/classification_evaluation",
+        "datamodule.num_workers=1",
+        "datamodule.batch_size=16",
+        "logger=csv",
+        "task.device=cpu",
+    ]
+
+    run_inference_experiments(
+        test_overrides=test_overrides,
+        data_path=data_path,
+        train_path=train_path,
+        test_path=test_path,
+        export_types=BASE_EXPORT_TYPES,
+    )
 
     shutil.rmtree(tmp_path)
 
 
-def test_train_multilabel_classification(
+def test_multilabel_classification(
     tmp_path: Path, base_multilabel_classification_dataset: base_multilabel_classification_dataset
 ):
+    """Test the training and evaluation of a torch based multilabel classification model."""
     data_path, arguments = base_multilabel_classification_dataset
 
     overrides = [
         "experiment=base/classification/multilabel_classification",
-        "trainer=lightning_cpu",
         "trainer.devices=1",
         f"datamodule.data_path={data_path}",
         f"datamodule.images_and_labels_file={Path(data_path) / 'samples.txt'}",
         f"model.classifier.out_features={len(arguments.samples)}",
         "backbone=resnet18",
         "trainer.max_epochs=1",
-    ] + BASE_EXPERIMENT_OVERRIDES
+        f"export.types=[{','.join(BASE_EXPORT_TYPES)}]",
+    ]
+    trainer_overrides = setup_trainer_for_lightning()
+    overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=tmp_path)
 
diff --git a/tests/tasks/test_segmentation.py b/tests/tasks/test_segmentation.py
index 75cb393d..af3ba6a8 100644
--- a/tests/tasks/test_segmentation.py
+++ b/tests/tasks/test_segmentation.py
@@ -2,14 +2,24 @@
 import os
 import shutil
 from pathlib import Path
+from typing import List
 
 import pytest
 
+from quadra.utils.export import get_export_extension
 from quadra.utils.tests.fixtures import base_binary_segmentation_dataset, base_multiclass_segmentation_dataset
-from quadra.utils.tests.helpers import execute_quadra_experiment
+from quadra.utils.tests.helpers import check_deployment_model, execute_quadra_experiment, setup_trainer_for_lightning
+
+try:
+    import onnx  # noqa
+    import onnxruntime  # noqa
+    import onnxsim  # noqa
+
+    ONNX_AVAILABLE = True
+except ImportError:
+    ONNX_AVAILABLE = False
 
 BASE_EXPERIMENT_OVERRIDES = [
-    "trainer=lightning_cpu",
     "trainer.devices=1",
     "trainer.max_epochs=1",
     "datamodule.num_workers=1",
@@ -20,10 +30,46 @@
     "logger=csv",
 ]
 
+BASE_EXPORT_TYPES = ["torchscript"] if not ONNX_AVAILABLE else ["torchscript", "onnx"]
+
+
+def _run_inference_experiment(
+    test_overrides: List[str], data_path: str, train_path: str, test_path: str, export_type: str
+):
+    """Run an inference experiment for the given export type."""
+    extension = get_export_extension(export_type)
+
+    test_overrides.append(f"datamodule.data_path={data_path}")
+    test_overrides.append(f"task.model_path={os.path.join(train_path, 'deployment_model', f'model.{extension}')}")
+
+    execute_quadra_experiment(overrides=test_overrides, experiment_path=test_path)
+
+
+def run_inference_experiments(
+    test_overrides: List[str], data_path: str, train_path: str, test_path: str, export_types: List[str]
+):
+    """Run inference experiments for the given export types."""
+    for export_type in export_types:
+        cwd = os.getcwd()
+        check_deployment_model(export_type=export_type)
+
+        _run_inference_experiment(
+            test_overrides=test_overrides,
+            data_path=data_path,
+            train_path=train_path,
+            test_path=test_path,
+            export_type=export_type,
+        )
+
+        # Change back to the original working directory
+        os.chdir(cwd)
+
 
 @pytest.mark.parametrize("generate_report", [True, False])
 def test_smp_binary(
-    tmp_path: Path, base_binary_segmentation_dataset: base_binary_segmentation_dataset, generate_report: bool
+    tmp_path: Path,
+    base_binary_segmentation_dataset: base_binary_segmentation_dataset,
+    generate_report: bool,
 ):
     data_path, _, _ = base_binary_segmentation_dataset
 
@@ -37,19 +83,26 @@ def test_smp_binary(
         f"datamodule.data_path={data_path}",
         f"task.report={generate_report}",
         "task.evaluate.analysis=false",
+        f"export.types=[{','.join(BASE_EXPORT_TYPES)}]",
     ]
+    trainer_overrides = setup_trainer_for_lightning()
     overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=train_path)
 
-    trained_model_path = os.path.join(train_path, "deployment_model/model.pt")
     inference_overrides = [
         "experiment=base/segmentation/smp_evaluation",
-        f"datamodule.data_path={data_path}",
-        f"task.model_path={trained_model_path}",
         "task.device=cpu",
     ] + BASE_EXPERIMENT_OVERRIDES
-    execute_quadra_experiment(overrides=inference_overrides, experiment_path=test_path)
+
+    run_inference_experiments(
+        test_overrides=inference_overrides,
+        data_path=data_path,
+        train_path=train_path,
+        test_path=test_path,
+        export_types=BASE_EXPORT_TYPES,
+    )
 
     shutil.rmtree(tmp_path)
 
@@ -72,20 +125,27 @@ def test_smp_multiclass(tmp_path: Path, base_multiclass_segmentation_dataset: ba
         f"datamodule.data_path={data_path}",
         f"datamodule.idx_to_class={idx_to_class_parameter}",
         "task.evaluate.analysis=false",
+        f"export.types=[{','.join(BASE_EXPORT_TYPES)}]",
     ]
+    trainer_overrides = setup_trainer_for_lightning()
     overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=train_path)
 
-    trained_model_path = os.path.join(train_path, "deployment_model/model.pt")
     inference_overrides = [
         "experiment=base/segmentation/smp_multiclass_evaluation",
-        f"datamodule.data_path={data_path}",
-        f"task.model_path={trained_model_path}",
         f"datamodule.idx_to_class={idx_to_class_parameter}",
         "task.device=cpu",
     ] + BASE_EXPERIMENT_OVERRIDES
-    execute_quadra_experiment(overrides=inference_overrides, experiment_path=test_path)
+
+    run_inference_experiments(
+        test_overrides=inference_overrides,
+        data_path=data_path,
+        train_path=train_path,
+        test_path=test_path,
+        export_types=BASE_EXPORT_TYPES,
+    )
 
     shutil.rmtree(tmp_path)
 
@@ -105,8 +165,11 @@ def test_smp_multiclass_with_binary_dataset(
         f"datamodule.data_path={data_path}",
         f"datamodule.idx_to_class={idx_to_class_parameter}",
         "task.evaluate.analysis=false",
+        f"export.types=[{','.join(BASE_EXPORT_TYPES)}]",
     ]
+    trainer_overrides = setup_trainer_for_lightning()
     overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=tmp_path)
 
diff --git a/tests/tasks/test_ssl.py b/tests/tasks/test_ssl.py
index 895f7f60..3b2fb4c9 100644
--- a/tests/tasks/test_ssl.py
+++ b/tests/tasks/test_ssl.py
@@ -3,10 +3,20 @@
 from pathlib import Path
 
 from quadra.utils.tests.fixtures import base_classification_dataset
-from quadra.utils.tests.helpers import execute_quadra_experiment
+from quadra.utils.tests.helpers import execute_quadra_experiment, setup_trainer_for_lightning
+
+try:
+    import onnx  # noqa
+    import onnxruntime  # noqa
+    import onnxsim  # noqa
+
+    ONNX_AVAILABLE = True
+except ImportError:
+    ONNX_AVAILABLE = False
+
+BASE_EXPORT_TYPES = ["torchscript"] if not ONNX_AVAILABLE else ["torchscript", "onnx"]
 
 BASE_EXPERIMENT_OVERRIDES = [
-    "trainer=lightning_cpu",
     "trainer.devices=1",
     "trainer.max_epochs=1",
     "trainer.check_val_every_n_epoch=1",
@@ -19,6 +29,7 @@
     "+trainer.limit_val_batches=1",
     "+trainer.limit_test_batches=1",
     "logger=csv",
+    f"export.types=[{','.join(BASE_EXPORT_TYPES)}]",
 ]
 
 
@@ -29,7 +40,10 @@ def test_simsiam(tmp_path: Path, base_classification_dataset: base_classificatio
         "experiment=base/ssl/simsiam",
         f"datamodule.data_path={data_path}",
         "backbone=resnet18",
-    ] + BASE_EXPERIMENT_OVERRIDES
+    ]
+    trainer_overrides = setup_trainer_for_lightning()
+    overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=tmp_path)
 
@@ -47,6 +61,7 @@ def test_dino(tmp_path: Path, base_classification_dataset: base_classification_d
         "loss.warmup_teacher_temp_epochs=1",
         "trainer.max_epochs=2",
     ]
+    overrides += setup_trainer_for_lightning()
 
     execute_quadra_experiment(overrides=overrides, experiment_path=tmp_path)
 
@@ -60,7 +75,10 @@ def test_simclr(tmp_path: Path, base_classification_dataset: base_classification
         "experiment=base/ssl/simclr",
         f"datamodule.data_path={data_path}",
         "backbone=resnet18",
-    ] + BASE_EXPERIMENT_OVERRIDES
+    ]
+    trainer_overrides = setup_trainer_for_lightning()
+    overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=tmp_path)
 
@@ -74,7 +92,10 @@ def test_byol(tmp_path: Path, base_classification_dataset: base_classification_d
         "experiment=base/ssl/byol",
         f"datamodule.data_path={data_path}",
         "backbone=resnet18",
-    ] + BASE_EXPERIMENT_OVERRIDES
+    ]
+    trainer_overrides = setup_trainer_for_lightning()
+    overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=tmp_path)
 
@@ -88,7 +109,10 @@ def test_barlow(tmp_path: Path, base_classification_dataset: base_classification
         "experiment=base/ssl/barlow",
         f"datamodule.data_path={data_path}",
         "backbone=resnet18",
-    ] + BASE_EXPERIMENT_OVERRIDES
+    ]
+    trainer_overrides = setup_trainer_for_lightning()
+    overrides += BASE_EXPERIMENT_OVERRIDES
+    overrides += trainer_overrides
 
     execute_quadra_experiment(overrides=overrides, experiment_path=tmp_path)