Repository: orobix/quadra (language: python)

Issue:
Feature/landing page
## Summary

This PR introduces some visual updates. Features are added:

- New Logo and theme
- Landing Page with animation

It also fixes the linking issues when the website deployed on custom `base_url` such as `quadra/`

## Type of Change

- Documentation update

## Checklist

Please confirm that the following tasks have been completed:

- [x] I have tested my changes locally and they work as expected. (Please describe the tests you performed.)
- [ ] I have added unit tests for my changes, or updated existing tests if necessary.
- [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/.github/workflows/docs.yaml b/.github/workflows/docs.yaml
new file mode 100644
index 00000000..27f1b0e5
--- /dev/null
+++ b/.github/workflows/docs.yaml
@@ -0,0 +1,47 @@
+name: Build and Deploy Documentation
+
+on:
+  push:
+    branches:
+      - main
+      - dev
+
+jobs:
+  build-and-deploy:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout Repository
+        uses: actions/checkout@v3
+
+      - name: Setup Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.9
+
+      - name: Install Dependencies
+        run: |
+          python -m pip install -U pip
+          python -m pip install -e ".[docs]" --no-cache-dir
+
+      - name: Determine Version
+        id: determine_version
+        run: |
+          if [[ "${{ github.ref }}" == "refs/heads/main" ]]; then
+            VERSION="v$(python -c "import quadra; print(quadra.__version__)")"
+          elif [[ "${{ github.ref }}" == "refs/heads/dev" ]]; then
+            VERSION="dev"
+          fi
+          echo "::set-output name=version::$VERSION"
+
+      - name: Build Documentation
+        run: |
+          git config user.name "${GITHUB_ACTOR}"
+          git config user.email "${GITHUB_ACTOR}@users.noreply.github.com"
+          git fetch origin gh-pages --depth=1
+          VERSION="${{ steps.determine_version.outputs.version }}"
+          if [[ "${{ github.ref }}" == "refs/heads/main" ]]; then
+            mike deploy --push --update-aliases $VERSION latest
+          else
+            mike deploy --push $VERSION
+          fi
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 03233dec..f7d371e4 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -5,8 +5,9 @@ repos:
     hooks:
       - id: ruff
         name: ruff
+        files: \.py$
         entry: ruff
-        language: system
+        language: python
         exclude: ^(tests|scripts|notebooks|quadra/utils/tests)/
   - repo: local
     hooks:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index af52a79f..d5254e05 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,11 +2,38 @@
 # Changelog
 All notable changes to this project will be documented in this file.
 
+### [1.1.0]
+
+#### Added
+
+- Add ModelManager class to manage model deployments on model tracking platforms (e.g. MLFlow).
+- Add automatic storage of file hashes in BaseDataModule class for better experiment reproducibility and tracking.
+- Automatically load transforms parameters from model info file in base Evaluation task.
+- Add support for vit explainability using Attention Gradient Rollout method.
+- Add export of pytorch model for Classification and SklearnClassification tasks.
+- Add automatical detection of model input shapes for better exportation capabilities. Add support for custom input shapes like models with multiple inputs.
+- Add documentation landing page, improve colore themes and logo.
+- Add github actions to automatically build and deploy documentation of main and dev PRs.
+
+#### Changed
+
+- Refactor evaluation task to be more generic and improve inheritance capabilities.
+- Refactor export_types parameter to export for better configurability of export parameters.
+- Change input_size model info parameter from HXW to a list of actual model parameters for inference (e.g [(3, 224, 224)]).
+
+#### Fixed
+
+- Fix gradcam not working properly with non rectangular images.
+- Fix logistic regression wrapper not working properly with 2 classes for torch classification.
+- Fix wrong typings in NetworkBuilder's init method.
+- Fix broken links in documentation.
+- Fix minor documentations issues.
+
 ### [1.0.2]
 
 #### Fixed
 
-- Fix anomaly detection training not working when images with non lowercase allowed extension are used (e.g. .BMP)
+- Fix anomaly detection training not working when images with non lowercase allowed extension are used (e.g. .BMP).
 
 #### Added
 
diff --git a/README.md b/README.md
index aa0d6898..9056ca9e 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,7 @@
-<h1 align="center">Quadra</h1>
+<h1></h1>
+<p align="center">
+  <img src="docs/images/quadra_text_logo.png" alt="Quadra Logo" width="100%">
+</p>
 
 <p align="center">
   <a href="https://orobix.github.io/quadra/latest/index.html">Docs</a> •
@@ -43,7 +46,7 @@ ______________________________________________________________________
 
 ## Quick Start Guide
 
-Currently we support installing from source since the library is not yet available on `PyPI`.
+Currently we support installing from source since the library is not yet available on `PyPI` and currently supported Python version is `3.9`.
 
 ```shell
 pip install git+https://github.com/orobix/quadra.git
@@ -109,7 +112,8 @@ This command sets the `MLFLOW_TRACKING_URI` variable to `http://localhost:5000`.
 
 **Adding it to your environment file:**
 
-You can also add the `MLFLOW_TRACKING_URI` variable to your environment file (e.g., `.env`). Open the file in a text editor and add the following line:
+`Quadra` uses environment variables to store credentials and other sensitive information. Thanks to [`python-dotenv`](https://pypi.org/project/python-dotenv/) library, you can create a `.env` file in the main folder of your project and store the credentials there. During the runtime, the library will automatically load the environment variables from the `.env` file. You can also add the `MLFLOW_TRACKING_URI` variable to your environment file (e.g., `.env`). Open the file in a text editor and add the following line:
+
 ```
 MLFLOW_TRACKING_URI=http://localhost:5000
 ```
@@ -218,7 +222,7 @@ quadra experiment=custom/<config_name> logger=csv
 
 It will run the experiment using the configuration file you have just created and it will apply the default parameters from the classification configuration file. Furthermore, it will log the metrics to a csv file. You can add or customize the parameters in the configuration file to fit your needs.
 
-For more information about advanced usage, please check [tutorials](/tutorials/configurations.html) and [task specific examples](/tutorials/examples/classification.html).
+For more information about advanced usage, please check [tutorials](https://orobix.github.io/quadra/latest/tutorials/configurations.html) and [task specific examples](https://orobix.github.io/quadra/latest/tutorials/examples/classification.html).
 
 ## Development
 
diff --git a/docs/css/colors.css b/docs/css/colors.css
new file mode 100644
index 00000000..68ee4fe2
--- /dev/null
+++ b/docs/css/colors.css
@@ -0,0 +1,13 @@
+:root {
+  --rolling-text-offset: 0.3em;
+  --quadra-blue: #0ecffe;
+  --quadra-blue-dark: #0699dd;
+  --quadra-green: #18e198;
+  --quadra-green-dark: #0d9548;
+
+  --quadra-purple: #8a7cfb;
+  --quadra-purple-dark: #633e9c;
+  --quadra-orange: #fa7671;
+  --quadra-orange-dark: #f45f7f;
+  --quadra-logo-single: #FFFFFF;
+}
\ No newline at end of file
diff --git a/docs/css/landing.css b/docs/css/landing.css
new file mode 100644
index 00000000..268a40bb
--- /dev/null
+++ b/docs/css/landing.css
@@ -0,0 +1,209 @@
+*, ::before, ::after {
+  --tw-border-spacing-x: 0;
+  --tw-border-spacing-y: 0;
+  --tw-translate-x: 0;
+  --tw-translate-y: 0;
+  --tw-rotate: 0;
+  --tw-skew-x: 0;
+  --tw-skew-y: 0;
+  --tw-scale-x: 1;
+  --tw-scale-y: 1;
+  --tw-pan-x:  ;
+  --tw-pan-y:  ;
+  --tw-pinch-zoom:  ;
+  --tw-scroll-snap-strictness: proximity;
+  --tw-gradient-from-position:  ;
+  --tw-gradient-via-position:  ;
+  --tw-gradient-to-position:  ;
+  --tw-ordinal:  ;
+  --tw-slashed-zero:  ;
+  --tw-numeric-figure:  ;
+  --tw-numeric-spacing:  ;
+  --tw-numeric-fraction:  ;
+  --tw-ring-inset:  ;
+  --tw-ring-offset-width: 0px;
+  --tw-ring-offset-color: #fff;
+  --tw-ring-color: rgb(59 130 246 / 0.5);
+  --tw-ring-offset-shadow: 0 0 #0000;
+  --tw-ring-shadow: 0 0 #0000;
+  --tw-shadow: 0 0 #0000;
+  --tw-shadow-colored: 0 0 #0000;
+  --tw-blur:  ;
+  --tw-brightness:  ;
+  --tw-contrast:  ;
+  --tw-grayscale:  ;
+  --tw-hue-rotate:  ;
+  --tw-invert:  ;
+  --tw-saturate:  ;
+  --tw-sepia:  ;
+  --tw-drop-shadow:  ;
+  --tw-backdrop-blur:  ;
+  --tw-backdrop-brightness:  ;
+  --tw-backdrop-contrast:  ;
+  --tw-backdrop-grayscale:  ;
+  --tw-backdrop-hue-rotate:  ;
+  --tw-backdrop-invert:  ;
+  --tw-backdrop-opacity:  ;
+  --tw-backdrop-saturate:  ;
+  --tw-backdrop-sepia:  
+}
+
+::backdrop {
+  --tw-border-spacing-x: 0;
+  --tw-border-spacing-y: 0;
+  --tw-translate-x: 0;
+  --tw-translate-y: 0;
+  --tw-rotate: 0;
+  --tw-skew-x: 0;
+  --tw-skew-y: 0;
+  --tw-scale-x: 1;
+  --tw-scale-y: 1;
+  --tw-pan-x:  ;
+  --tw-pan-y:  ;
+  --tw-pinch-zoom:  ;
+  --tw-scroll-snap-strictness: proximity;
+  --tw-gradient-from-position:  ;
+  --tw-gradient-via-position:  ;
+  --tw-gradient-to-position:  ;
+  --tw-ordinal:  ;
+  --tw-slashed-zero:  ;
+  --tw-numeric-figure:  ;
+  --tw-numeric-spacing:  ;
+  --tw-numeric-fraction:  ;
+  --tw-ring-inset:  ;
+  --tw-ring-offset-width: 0px;
+  --tw-ring-offset-color: #fff;
+  --tw-ring-color: rgb(59 130 246 / 0.5);
+  --tw-ring-offset-shadow: 0 0 #0000;
+  --tw-ring-shadow: 0 0 #0000;
+  --tw-shadow: 0 0 #0000;
+  --tw-shadow-colored: 0 0 #0000;
+  --tw-blur:  ;
+  --tw-brightness:  ;
+  --tw-contrast:  ;
+  --tw-grayscale:  ;
+  --tw-hue-rotate:  ;
+  --tw-invert:  ;
+  --tw-saturate:  ;
+  --tw-sepia:  ;
+  --tw-drop-shadow:  ;
+  --tw-backdrop-blur:  ;
+  --tw-backdrop-brightness:  ;
+  --tw-backdrop-contrast:  ;
+  --tw-backdrop-grayscale:  ;
+  --tw-backdrop-hue-rotate:  ;
+  --tw-backdrop-invert:  ;
+  --tw-backdrop-opacity:  ;
+  --tw-backdrop-saturate:  ;
+  --tw-backdrop-sepia:  
+}
+
+.static {
+  position: static
+}
+
+.absolute {
+  position: absolute
+}
+
+.relative {
+  position: relative
+}
+
+.mb-12 {
+  margin-bottom: 3rem
+}
+
+.mt-2 {
+  margin-top: 0.5rem
+}
+
+.block {
+  display: block
+}
+
+.flex {
+  display: flex
+}
+
+.table {
+  display: table
+}
+
+.contents {
+  display: contents
+}
+
+.hidden {
+  display: none
+}
+
+.transform {
+  transform: translate(var(--tw-translate-x), var(--tw-translate-y)) rotate(var(--tw-rotate)) skewX(var(--tw-skew-x)) skewY(var(--tw-skew-y)) scaleX(var(--tw-scale-x)) scaleY(var(--tw-scale-y))
+}
+
+.px-6 {
+  padding-left: 1.5rem;
+  padding-right: 1.5rem
+}
+
+.py-24 {
+  padding-top: 6rem;
+  padding-bottom: 6rem
+}
+
+.text-center {
+  text-align: center
+}
+
+.text-2xl {
+  font-size: 1.5rem;
+  line-height: 2rem
+}
+
+.text-base {
+  font-size: 1rem;
+  line-height: 1.5rem
+}
+
+.font-bold {
+  font-weight: 700
+}
+
+.tracking-tight {
+  letter-spacing: -0.025em
+}
+
+.transition {
+  transition-property: color, background-color, border-color, text-decoration-color, fill, stroke, opacity, box-shadow, transform, filter, -webkit-backdrop-filter;
+  transition-property: color, background-color, border-color, text-decoration-color, fill, stroke, opacity, box-shadow, transform, filter, backdrop-filter;
+  transition-property: color, background-color, border-color, text-decoration-color, fill, stroke, opacity, box-shadow, transform, filter, backdrop-filter, -webkit-backdrop-filter;
+  transition-timing-function: cubic-bezier(0.4, 0, 0.2, 1);
+  transition-duration: 150ms
+}
+
+@media (min-width: 640px) {
+  .sm\:text-2xl {
+    font-size: 1.5rem;
+    line-height: 2rem
+  }
+
+  .sm\:text-sm {
+    font-size: 0.875rem;
+    line-height: 1.25rem
+  }
+}
+
+@media (min-width: 768px) {
+  .md\:text-5xl {
+    font-size: 3rem;
+    line-height: 1
+  }
+}
+
+@media (min-width: 1280px) {
+  .xl\:text-6xl {
+    font-size: 3.75rem;
+    line-height: 1
+  }
+}
\ No newline at end of file
diff --git a/docs/css/mkdocstrings.css b/docs/css/mkdocstrings.css
index e5260c57..831a9a77 100644
--- a/docs/css/mkdocstrings.css
+++ b/docs/css/mkdocstrings.css
@@ -4,17 +4,17 @@ div.doc-contents:not(.first) {
   border-left: 4px solid rgba(230, 230, 230);
   margin-bottom: 80px;
 }
-  
-  /* Avoid breaking parameters name, etc. in table cells.
+
+/* Avoid breaking parameters name, etc. in table cells.
 td code {
     word-break: normal !important;
   } */
 
 .md-sidebar__inner {
-    border-right: .05rem solid rgba(0,0,0,.07);
+  border-right: .05rem solid rgba(0, 0, 0, .07);
 }
 
 .md-nav--secondary {
-    border-left: .15rem solid #ef5552;
-    border-right: none;
+  border-left: .15rem solid var(--md-accent-fg-color);
+  border-right: none;
 }
\ No newline at end of file
diff --git a/docs/gen_ref_pages.py b/docs/gen_ref_pages.py
index dd91dd81..f06ea3c9 100644
--- a/docs/gen_ref_pages.py
+++ b/docs/gen_ref_pages.py
@@ -80,7 +80,7 @@ def add_submodules_as_list(parent_folder: Path) -> str:
     readme = read_fd.read()
     readme = readme.replace("# Quadra: Deep Learning Experiment Orchestration Library", "# Home")
     readme = readme.replace("docs/", "")
-    with mkdocs_gen_files.open("index.md", "w") as nav_file:  # (2)
+    with mkdocs_gen_files.open("getting_started.md", "w") as nav_file:  # (2)
         nav_file.write(readme)
 
 with open("CHANGELOG.md", "r") as change_fd:
diff --git a/docs/images/favicon.png b/docs/images/favicon.png
index 58592639..765b55d3 100644
Binary files a/docs/images/favicon.png and b/docs/images/favicon.png differ
diff --git a/docs/images/model_manager/model_compare.png b/docs/images/model_manager/model_compare.png
new file mode 100644
index 00000000..6cf4eecf
Binary files /dev/null and b/docs/images/model_manager/model_compare.png differ
diff --git a/docs/images/model_manager/model_description.png b/docs/images/model_manager/model_description.png
new file mode 100644
index 00000000..16af05d3
Binary files /dev/null and b/docs/images/model_manager/model_description.png differ
diff --git a/docs/images/model_manager/registered_model.png b/docs/images/model_manager/registered_model.png
new file mode 100644
index 00000000..766b5b33
Binary files /dev/null and b/docs/images/model_manager/registered_model.png differ
diff --git a/docs/images/model_manager/stage_model.png b/docs/images/model_manager/stage_model.png
new file mode 100644
index 00000000..392658da
Binary files /dev/null and b/docs/images/model_manager/stage_model.png differ
diff --git a/docs/images/quadra_logo.svg b/docs/images/quadra_logo.svg
new file mode 100644
index 00000000..7ce01b9e
--- /dev/null
+++ b/docs/images/quadra_logo.svg
@@ -0,0 +1,17 @@
+<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" version="1.1" id="Layer_1" x="0px" y="0px" style="enable-background:new 0 0 1920 1080;" xml:space="preserve" viewBox="781.3 406.3 286.6 311.4">
+<style type="text/css">
+	.q-logo-single{fill:#FFFFFF;}
+</style>
+<g>
+	<g>
+		<path class="q-logo-single" d="M781.3,548.3c0-81.8,63.5-142,143.3-142s143.3,60.2,143.3,142c0,81.8-63.5,142-143.3,142    S781.3,630.1,781.3,548.3z M1009.2,548.3c0-51.7-34.2-91.2-84.6-91.2c-50.5,0-84.6,39.5-84.6,91.2s34.2,91.2,84.6,91.2    C975,639.4,1009.2,600,1009.2,548.3z"/>
+	</g>
+	<g>
+		<g>
+			<path class="q-logo-single" d="M997.2,655.6l-47.9-34.9c-14.4,19.8-10,47.7,9.8,62.1l47.9,34.9C1021.3,697.9,1016.9,670,997.2,655.6z"/>
+			<path class="q-logo-single" d="M994.7,640.1l-47.9-34.9c-14.4,19.8-10,47.7,9.8,62.1l47.9,34.9C1018.9,682.4,1014.5,654.5,994.7,640.1z"/>
+			<path class="q-logo-single" d="M992.3,624.6l-47.9-34.9c-14.4,19.8-10,47.7,9.8,62.1l47.9,34.9C1016.5,666.9,1012.1,639,992.3,624.6z"/>
+		</g>
+	</g>
+</g>
+</svg>
\ No newline at end of file
diff --git a/docs/images/quadra_text_logo.png b/docs/images/quadra_text_logo.png
new file mode 100644
index 00000000..9f2c2c2d
Binary files /dev/null and b/docs/images/quadra_text_logo.png differ
diff --git a/docs/index.md b/docs/index.md
index b2b17e35..c6689b08 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -1,4 +1,5 @@
-# Home
-
-Welcome to `quadra` library documentation. Please visit [tutorials](/tutorials/install.html) to get familiar with the library.
+---
+template: home.html
+title: Home
+---
 
diff --git a/docs/tailwind.config.js b/docs/tailwind.config.js
new file mode 100644
index 00000000..57fefb40
--- /dev/null
+++ b/docs/tailwind.config.js
@@ -0,0 +1,17 @@
+/** @type {import('tailwindcss').Config} */
+module.exports = {
+  content: ["./docs_overrides/home.html"],
+  // darkMode: "class",
+  theme: {
+    fontFamily: {
+      sans: ["Roboto", "sans-serif"],
+      body: ["Roboto", "sans-serif"],
+      mono: ["ui-monospace", "monospace"],
+    },
+  },
+  corePlugins: {
+    preflight: false,
+  },
+  plugins: [],
+}
+
diff --git a/docs/template/landing.css b/docs/template/landing.css
new file mode 100644
index 00000000..bd6213e1
--- /dev/null
+++ b/docs/template/landing.css
@@ -0,0 +1,3 @@
+@tailwind base;
+@tailwind components;
+@tailwind utilities;
\ No newline at end of file
diff --git a/docs/tutorials/configurations.md b/docs/tutorials/configurations.md
index 45e6dd94..bd442ff3 100644
--- a/docs/tutorials/configurations.md
+++ b/docs/tutorials/configurations.md
@@ -79,7 +79,7 @@ For example we can set the seed, decide for a run name, or set the log level.
 
 ### Datamodule
 
-Datamodule setting files are used to configure the datamodule which manages the datasets and dataloaders used in experiment. For a detailed explanatation on how to implement `DataModule` classes, please refer to the [datamodule documentation](/tutorials/datamodules.html).
+Datamodule setting files are used to configure the datamodule which manages the datasets and dataloaders used in experiment. For a detailed explanation on how to implement `DataModule` classes, please refer to the [datamodule documentation](../tutorials/datamodules.md).
 
 Here is the structure of the folder:
 
@@ -175,10 +175,13 @@ model:
 task:
   lr_multiplier: 0.0
   run_test: True
-  export_type: [torchscript]
   report: True
   output:
     example: True
+  export_config:
+    types: [torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
 
 core:
   tag: "run"
@@ -327,7 +330,7 @@ Each scheduler file defines how we initialize the learning rate schedulers with
 
 ### Task
 
-The tasks are the building blocks containing the actual training and evaluation logic. They are discussed in more details in the [tasks](/tutorials/tasks.html) section.
+The tasks are the building blocks containing the actual training and evaluation logic. They are discussed in more details in the [tasks](../tutorials/tasks.md) section.
 
 ### Trainer
 
diff --git a/docs/tutorials/contribution.md b/docs/tutorials/contribution.md
index f6e4eb91..fdf2d4a9 100644
--- a/docs/tutorials/contribution.md
+++ b/docs/tutorials/contribution.md
@@ -9,7 +9,15 @@ In this guide, we'll cover the steps you should take to ensure your code meets t
 
 ## Setting up the Environment
 
-Before contributing to the repository, you'll need to set up your development environment. Please check the [installation guide](/tutorials/install.html) for instructions on how to set up your environment.
+Before contributing to the repository, you'll need to set up your development environment. Please check the [Getting Started Guide](../getting_started.md) for instructions on how to set up your environment.
+
+After setting up your environment you can install `Quadra` Library in different ways:
+
+!!!info
+
+    - `pip install -e .[dev]` (for development) 
+    - `pip install -e .[docs]` (for documentation)
+    - `pip install -e .[test]` (for testing)
 
 ## Pre-commit Hooks
 
diff --git a/docs/tutorials/devices_setup.md b/docs/tutorials/devices_setup.md
index dbcb7956..54556e65 100644
--- a/docs/tutorials/devices_setup.md
+++ b/docs/tutorials/devices_setup.md
@@ -36,13 +36,15 @@ For Sklearn based tasks there's generally a `device` field in the configuration
 ```yaml
 _target_: quadra.tasks.SklearnClassification
 device: "cuda:0"
-export_type: [torchscript]
 output:
   folder: "classification_experiment"
   save_backbone: false
   report: true
   example: true
   test_full_data: true
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
 ```
 
 You can change the device to `cpu` or a different cuda device depending on your needs.
diff --git a/docs/tutorials/documentation.md b/docs/tutorials/documentation.md
index 09a1605a..9f7d6403 100644
--- a/docs/tutorials/documentation.md
+++ b/docs/tutorials/documentation.md
@@ -33,6 +33,14 @@ nav:
   # ...
 ```
 
+## Developing Landing Page
+
+
+Landing page extends the `home.html` from `mkdocs-material` theme. It uses `tailwindcss` for styling. The developer has to run the following command to update the `tailwindcss` styles.
+
+```bash
+npx tailwindcss -i ./docs/template/landing.css -o ./docs/css/landing.css --watch
+```
 
 ## Features
 
diff --git a/docs/tutorials/examples/anomaly_detection.md b/docs/tutorials/examples/anomaly_detection.md
index b742cce2..d3788f76 100644
--- a/docs/tutorials/examples/anomaly_detection.md
+++ b/docs/tutorials/examples/anomaly_detection.md
@@ -10,7 +10,7 @@ This example will demonstrate how to create custom experiments starting from def
 
 Let's start with the dataset that we are going to use. Since we are using the base anomaly datamodule, images
 and masks must be arranged in a folder structure that follows the anomaly datamodule guidelines defined in the 
-[anomaly datamodule documentation](/tutorials/datamodules.html#anomaly-detection). 
+[anomaly datamodule documentation](../../tutorials/datamodules.md#anomaly-detection). 
 For this example, we will use the `mnist` dataset (using 9 as good and all the other number as anomalies), the dataset will be automatically downloaded by the generic experiment described next.
 
 ```tree
diff --git a/docs/tutorials/examples/classification.md b/docs/tutorials/examples/classification.md
index 2beafaf7..4c21705a 100644
--- a/docs/tutorials/examples/classification.md
+++ b/docs/tutorials/examples/classification.md
@@ -68,6 +68,16 @@ We give a small description for each tweakable parameter inside the base datamod
 - `label_map`: You can map classes to other ones in order to group more sub-classes into one macro-class.
 - `class_to_idx`: Map classes to indexes, if you need to ensure that one specific mapping is respected (otherwise a ordered class_to_idx is built)
 
+If you want to choose which files to put in training/val/test sets, provide the split files' paths in the respective fields in the datamodule config (train_split_file, val_split_file, test_split_file). You can leave val_split_file empty and the training set will be automatically split in train/val.
+
+The content of the files must be formatted in the same way for `train.txt`/`val.txt`/`test.txt`:
+
+```txt
+images/abc.xyz,class_1,class_2
+images/abc_2.xyz,class_3
+```
+
+The first column is the path to the image, while the other columns are the labels associated with that image. The labels must be separated by a comma.
 
 ### Experiment
 
@@ -94,17 +104,20 @@ print_config: true
 
 model:
   num_classes: ???
-  gradcam: true
   module:
     lr_scheduler_interval: "epoch"
 
 task:
   lr_multiplier: 0.0
+  gradcam: true
   run_test: True
-  export_type: [torchscript]
   report: True
   output:
     example: True
+  export_config:
+    types: [pytorch, torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
 
 core:
   tag: "run"
@@ -169,23 +182,30 @@ datamodule:
     class_3: 2
 
 task:
+  gradcam: True # Enable gradcam computation during evaluation
   run_test: True # Perform test evaluation at the end of training
-  export_type: [torchscript]
   report: True 
   output:
     example: True # Generate an example of concordants and discordants predictions for each class
+  export_config:
+    types: [pytorch, torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
    
 model:
   num_classes: 3 # This is very important
   module:
     lr_scheduler_interval: "epoch"
-    gradcam: True # Enable gradcam computation during evaluation
     
 backbone:
   model:
     pretrained: True
     freeze: False
-  freeze_parameters_name: # Here we could specify a list of layer names to freeze
+  freeze_parameters_name:
+    - conv1
+    - bn1
+    - layer1
+    - layer2
 
 core:
   tag: "run"
@@ -258,8 +278,14 @@ core:
   upload_artifacts: true
   name: classification_evalutation_base
 
+logger:
+  mlflow:
+    experiment_name: name_of_the_experiment
+    run_name: ${core.name}
+
 task:
   _target_: quadra.tasks.ClassificationEvaluation
+  gradcam: true
   output:
     example: true
   model_path: ???
@@ -283,9 +309,6 @@ datamodule:
     class_2: 1
     class_3: 2
 
-model:
-  num_classes: 3
-
 core:
   tag: "run"
   upload_artifacts: true
@@ -294,7 +317,7 @@ core:
 task:
   output:
       example: true
-  model_path: path/to/deployment_model
+  model_path: path/to/model.pth
 ```
 
 Notice that we must provide the path to a deployment model file that will be used to perform inferences. In this case class_to_idx is mandatory (we can not infer it from a test-set). We suggest to be careful to set the same class_to_idx that has been used to train the model.
diff --git a/docs/tutorials/examples/multilabel_classification.md b/docs/tutorials/examples/multilabel_classification.md
index 9b2eeedb..cd9fe0b6 100644
--- a/docs/tutorials/examples/multilabel_classification.md
+++ b/docs/tutorials/examples/multilabel_classification.md
@@ -136,10 +136,13 @@ model:
 
 task:
   run_test: True # Perform test evaluation at the end of training
-  export_type: [torchscript]
   report: False 
   output:
     example: False 
+  export_config:
+    types: [torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
   
 logger:
   mlflow:
diff --git a/docs/tutorials/examples/segmentation.md b/docs/tutorials/examples/segmentation.md
index 280e9041..5aa8fb43 100644
--- a/docs/tutorials/examples/segmentation.md
+++ b/docs/tutorials/examples/segmentation.md
@@ -7,7 +7,7 @@ In this page, we will show you how to run a segmentation experiment (either bina
 This example will demonstrate how to create a custom experiment starting from default settings.
 ### Dataset
 
-Let's start with the dataset we are going to use. Since we are using base segmentation datamodule, we must arrange our images and masks in a folder structure that follows the segmentation datamodule guideline defined in the [segmentation datamodule documentation](/tutorials/datamodules.html#segmentation).
+Let's start with the dataset we are going to use. Since we are using base segmentation datamodule, we must arrange our images and masks in a folder structure that follows the segmentation datamodule guideline defined in the [segmentation datamodule documentation](../../tutorials/datamodules.md#segmentation).
 Imagine that we have a dataset with the following structure:
 
 ```tree
@@ -167,11 +167,14 @@ backbone:
     classes: 4 # The total number of classes (background + foreground)
 
 task:
-  export_type: [torchscript] # Export model to torchscript for inference
   run_test: true # run test after training is completed
   report: false # allows to generate reports
   evaluate: # custom evaluation toggles
     analysis: false # Perform in depth analysis
+  export_config:
+    types: [torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
     
 datamodule:
   data_path: /path/to/the/dataset # change the path to the dataset
@@ -197,9 +200,9 @@ 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.
 
-By default data will be logged to mlflow if it is [configured properly](/tutorials/install.html#mlflow-credentials). If mlflow is not available it's possible to configure a simple csv logger by adding an override to the file above:
+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:
   
 ```yaml
 # @package _global_
diff --git a/docs/tutorials/examples/sklearn_classification.md b/docs/tutorials/examples/sklearn_classification.md
index 6df2a5da..7c14bc26 100644
--- a/docs/tutorials/examples/sklearn_classification.md
+++ b/docs/tutorials/examples/sklearn_classification.md
@@ -100,6 +100,12 @@ backbone:
     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"
@@ -112,6 +118,7 @@ datamodule:
 ```
 
 By default the experiment will use dino_vitb8 as backbone, resizing the images to 224x224 and training a logistic regression classifier. Setting the `n_splits` parameter to 1 will use a standard 70/30 train/validation split (given the parameters specified in the base datamodule) instead of cross validation.
+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
@@ -139,16 +146,19 @@ datamodule:
 
 task:
   device: cuda:0
-  export_type: [torchscript]
   output:
     folder: classification_experiment
     save_backbone: true
     report: true
     example: true
     test_full_data: true
+  export_config:
+    types: [pytorch, torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
 ```
 
-This will train a logistic regression classifier using a resnet18 backbone, resizing the images to 224x224 and using a 5-fold cross validation. The `class_to_idx` parameter is used to map the class names to indexes, the indexes will be used to train the classifier. The `output` parameter is used to specify the output folder and the type of output to save. The `export_type` parameter can be used to export the model in different formats, at the moment only `torchscript` is supported.
+This will train a logistic regression classifier using a resnet18 backbone, resizing the images to 224x224 and using a 5-fold cross validation. The `class_to_idx` parameter is used to map the class names to indexes, the indexes will be used to train the classifier. The `output` parameter is used to specify the output folder and the type of output to save. The `export_config.types` parameter can be used to export the model in different formats, at the moment `torchscript` and `pytorch` are supported.
 Since `save_backbone` is set to true, the backbone (in torchscript format) will be saved along with the classifier. `test_full_data` is used to specify if a final test should be performed on all the data (after training on the training and validation datasets).
 
 ### Run
@@ -186,15 +196,9 @@ The default experiment config is found under `configs/experiment/base/classifica
 defaults:
   - override /transforms: default_resize
   - override /task: sklearn_classification_test
-  - override /backbone: dino_vitb8
   - override /trainer: sklearn_classification
   - override /datamodule: base/sklearn_classification
 
-backbone:
-  model:
-    pretrained: true
-    freeze: true
-
 core:
   tag: run
   name: sklearn-classification-test
@@ -211,7 +215,6 @@ An actual configuration file based on the above could be this one (suppose it's
 # @package _global_
 defaults:
   - base/classification/sklearn_classification_test
-  - override /backbone: resnet18
   - _self_
 
 core:
@@ -223,14 +226,16 @@ datamodule:
 
 task:
   device: cuda:0
+  gradcam: true
   output:
     folder: classification_test_experiment
     report: true
     example: true
-  experiment_path:
+  model_path: ???
 ```
 
-This will test the model trained in the given experiment on the given dataset. The experiment results will be saved under the `classification_test_experiment` folder.
+This will test the model trained in the given experiment on the given dataset. The experiment results will be saved under the `classification_test_experiment` folder. If gradcam is set to True, original and gradcam results will be saved during the generate_report().
+Model_path must point to a model file. It could either be a '.pt'/'.pth' or a backbone_config '.yaml' file.
 
 ### Run
 
diff --git a/docs/tutorials/examples/sklearn_patch_classification.md b/docs/tutorials/examples/sklearn_patch_classification.md
index 2c185b0c..35db4453 100644
--- a/docs/tutorials/examples/sklearn_patch_classification.md
+++ b/docs/tutorials/examples/sklearn_patch_classification.md
@@ -222,11 +222,14 @@ task:
   device: cuda:2
   output:
     folder: classification_patch_experiment
-    export_type: [torchscript]
     save_backbone: false
     report: true
     example: true
     reconstruction_method: major_voting
+  export_config:
+    types: [torchscript]
+    input_shapes: # Redefine the input shape if not automatically inferred
+
 ```
 
 This will train a resnet18 model on the given dataset, using 256 as batch size and skipping the background class during training.
@@ -265,15 +268,9 @@ The default experiment config is found under `configs/experiment/base/classifica
 defaults:
   - override /transforms: default_resize
   - override /task: sklearn_classification_patch_test
-  - override /backbone: dino_vitb8
   - override /datamodule: base/sklearn_classification_patch
   - override /trainer: sklearn_classification
 
-backbone:
-  model:
-    pretrained: true
-    freeze: true
-
 core:
   tag: "run"
   name: "sklearn-classification-patch-test"
@@ -307,11 +304,11 @@ task:
     report: true
     example: true
     reconstruction_method: major_voting
-  experiment_path: path_to_training_experiment
+  model_path: ???
 ```
 
 This will test the model trained in the given experiment on the given dataset. The experiment results will be saved under the `classification_patch_test` folder.
-Patch reconstruction will be performed using the `major_voting` method (can be this or `priority`). The `experiment_path` parameter is required to specify the path to the training experiment. In this case is not necessary to specify `class_to_idx` and `class_to_skip_training` as they will be loaded from the training experiment.
+Patch reconstruction will be performed using the `major_voting` method (can be this or `priority`). The `model_path` parameter is required to specify the path to the trained model. It could either be a '.pt'/'.pth' or a backbone_config '.yaml' file. In this case is not necessary to specify `class_to_idx` and `class_to_skip_training` as they will be loaded from the training experiment.
 
 ### Run
 
diff --git a/docs/tutorials/examples/ssl.md b/docs/tutorials/examples/ssl.md
index a2905821..836796e5 100644
--- a/docs/tutorials/examples/ssl.md
+++ b/docs/tutorials/examples/ssl.md
@@ -160,7 +160,7 @@ The output folder should contain the following entries:
 checkpoints  config_resolved.yaml  config_tree.txt  data  deployment_model  main.log
 ```
 
-The `checkpoints` folder contains the saved `pytorch` lightning checkpoints. The `data` folder contains a joblib version of the datamodule containing all parameters and dataset spits. The `deployment_model` folder contains the model ready for production in the format specified in the task `export_type` parameter (default `torchscript`). 
+The `checkpoints` folder contains the saved `pytorch` lightning checkpoints. The `data` folder contains a joblib version of the datamodule containing all parameters and dataset spits. The `deployment_model` folder contains the model ready for production in the format specified in the task `export_config.types` parameter (default `torchscript`). 
 
 ### Run (Advanced) - Changing transformations
 
diff --git a/docs/tutorials/install.md b/docs/tutorials/install.md
deleted file mode 100644
index 125e5a8b..00000000
--- a/docs/tutorials/install.md
+++ /dev/null
@@ -1,55 +0,0 @@
-# Installation
-
-## Virtual environment
-
-To start a new Python environment, run the following command:
-
-```bash
-cd `<your-project-directory>`
-python3 -m venv ./<environment-name>
-source ./<environment-nam>/bin/activate
-```
-
-If you prefer to use `conda` instead of `venv`, run the following command:
-
-```bash
-conda create -n <environment-name> python=3.9
-conda activate <environment-name>
-```
-
-!!!note
-    This library relies on python version 3.8 or 3.9 at the moment.
-
-## Installing the package
-
-Once the environment is activated, you can install the package using the following command:
-
-```bash
-pip install .
-```
-
-!!!note
-    This installs the package for deployment, there are other ways to install the package based on your needs:
-
-    - `pip install -e .[dev]` (for development) 
-    - `pip install -e .[docs]` (for documentation)
-    - `pip install -e .[test]` (for testing)
-
-## Usage
-
-When you install the package, you can use the following CLI command to run the experiments:
-
-```bash
-quadra <config-composition>
-```
-
-For example this is how a simple classification experiments on the `imagenette` dataset is run:
-
-```bash
-quadra experiment=generic/imagenette/classification/classification
-```
-
-## Credentials and Environment variables
-
-`quadra` uses environment variables to store credentials and other sensitive information. Thanks to [`python-dotenv`](https://pypi.org/project/python-dotenv/) library, you can create a `.env` file in the main folder of your project and store the credentials there. During the runtime, the library will automatically load the environment variables from the `.env` file.
-
diff --git a/docs/tutorials/model_management.md b/docs/tutorials/model_management.md
new file mode 100644
index 00000000..50b514be
--- /dev/null
+++ b/docs/tutorials/model_management.md
@@ -0,0 +1,150 @@
+# Model Management
+
+## Overview
+
+Model Management is the process of maintaining and orchestrating the lifecycle of machine learning models. It involves steps such as model creation, training, evaluation, deployment, and monitoring, among other steps. Model Versioning, on the other hand, allows data scientists to track and manage different versions of models during the model's lifecycle. It's crucial for reproducibility, collaboration, and consistent model performance. Quadra, out-of-the-box, supports managing models and their versions through [`MlflowModelManager`][quadra.utils.model_manager.MlflowModelManager] wrapper class with the help of [Mlflow](https://mlflow.org/) library. 
+
+ [`MlflowModelManager`][quadra.utils.model_manager.MlflowModelManager] is an extension of [`AbstractModelManager`][quadra.utils.model_manager.AbstractModelManager] which serves as a blueprint for model managers and specifies the required methods for managing models. It includes methods such as:
+
+- Register the model
+- Retrieve latest version
+- Transition the model to a new stage
+- Delete the model 
+
+By defining an abstract class, we establish a common interface that can be implemented by different model managers besides Mlflow.
+
+
+## Example Usage
+
+In this section, we will create example project for segmentation task and use [`MlflowModelManager`][quadra.utils.model_manager.MlflowModelManager] to manage the production model. `Quadra` provides a toy example where you can train a segmentation model for Oxford-IIIT Pet Dataset. 
+
+!!!note
+    You can find a detailed explanation for customizing the segmentation task under [Segmentation Example](/tutorials/examples/segmentation.md) section.
+
+First of all, we need to run `Mlflow` server with artifact store. You can find the instructions for running `Mlflow` server [here](https://mlflow.org/docs/latest/tracking.html#mlflow-tracking-servers). Let's open a new terminal and run the following command:
+
+```bash
+mlflow server \
+--backend-store-uri sqlite:///mlflow.db \
+--default-artifact-root file:///tmp/mlflow \
+--host 0.0.0.0
+```
+
+Then, we can start training from different terminal window while `Mlflow` server is running:
+
+```bash
+MLFLOW_TRACKING_URI="http://localhost:5000" \
+quadra experiment=generic/oxford_pet/segmentation/smp \
+trainer.max_epochs=5 \
+core.name=cats_vs_dogs \
+backbone.model.arch=unet,unetplusplus \
+backbone.model.encoder_name=resnet18,resnet50 \
+--multirun
+```
+
+This command will train a segmentation model for 5 epochs and save experiments run under `cats_vs_dogs` experiment tab. It will run the same experiment with different backbones or segmentation model architectures. After all trainings are completed, we open up the `Mlflow` UI and see experiments run under `cats_vs_dogs` directory. Under the artifacts section of each run, we can find `Mlflow` model artifact under the `deployment_model` directory. Other folders contains metadata or reports about the experiment run.
+
+### Registering the Model
+
+After all runs are completed, we can register the best model with the following command:
+
+```python
+import os
+from quadra.utils.model_manager import MlflowModelManager
+os.environ["MLFLOW_TRACKING_URI"] = "http://localhost:5000"
+
+manager = MlflowModelManager()
+manager.register_best_model(experiment_name="cats_vs_dogs",
+                            # metric we want to use for selecting the best model
+                            metric="val_loss_epoch", 
+                            # model path under the artifact store
+                            model_path="deployment_model", 
+                            # the name of the model to be registered
+                            model_name="cvsd_model", 
+                            # optional tags
+                            tags={"type":"segmentation"}, 
+                            # metric sorting order for selecting the best model
+                            mode="min")  
+```
+
+<p align="center">
+    <figure>
+        <img src="../images/model_manager/registered_model.png" title="Registered Model">
+        <figcaption>Registered Model</figcaption>
+    </figure>
+</p>
+
+
+### Staging the Model
+
+After registering the model, we can transition the model to a new stage. We can transition the model to `staging` stage with the following command:
+
+```python
+manager.transition_model(model_name="cvsd_model",
+                        version=1,
+                        stage="staging",
+                        description="Staging Model for demo")
+```
+
+<p align="center">
+    <figure>
+        <img src="../images/model_manager/stage_model.png" title="Staged Model">
+        <figcaption>Staged Model</figcaption>
+    </figure>
+</p>
+
+
+### Update the Model
+
+Let's say we want to update the model with better performance. Let's train the one of the configuration with more epochs and register the model again:
+
+```bash
+MLFLOW_TRACKING_URI="http://localhost:5000" \
+quadra experiment=generic/oxford_pet/segmentation/smp \
+trainer.max_epochs=20 \
+core.name=cats_vs_dogs \
+backbone.model.arch=unetplusplus \
+backbone.model.encoder_name=resnet50 
+```
+
+With this new model, we have improved `val_loss_epoch` metric compared to the previous model.
+
+<p align="center">
+    <figure>
+        <img src="../images/model_manager/model_compare.png" title="Comparison of Runs">
+        <figcaption>Comparison of Runs</figcaption>
+    </figure>
+</p>
+
+Let's register the new model with the following command:
+
+```python
+manager.register_model(model_location="runs:/<run-id>/deployment_model/model.pt",
+                       model_name="cvsd_model",
+                       tags={"type":"smp"},
+                       description="Better model"
+                       )
+```
+
+and finally transition the model to `production` stage:
+
+```python
+manager.transition_model(model_name="cvsd_model",
+                        version=2,
+                        stage="production",
+                        description="Production Model for demo")
+```
+
+!!!note
+    When we visit the model page from Mlflow UI, we can see the model version and stage information under the `Versions` tab. Moreover, we can see the model history automatically stored under the `Description` field.
+
+    <p align="center">
+        <figure>
+            <img src="../images/model_manager/model_description.png" title="Model Transition History">
+            <figcaption>Model Transition History</figcaption>
+        </figure>
+    </p>
+
+
+
+
diff --git a/docs/tutorials/tasks.md b/docs/tutorials/tasks.md
index 76d836e4..c05a3547 100644
--- a/docs/tutorials/tasks.md
+++ b/docs/tutorials/tasks.md
@@ -32,6 +32,7 @@ graph LR;
     L-->SG(Segmentation);
     L-->CL(Classification);
     T-->SCL(SklearnClassification);
+    T-->PSCL(PatchSklearnClassification);
     L-->SSL(SSL);
     L-->AD(AnomalibDetection)
 ```
@@ -42,10 +43,11 @@ graph LR;
 - **[`Segmentation Task`][quadra.tasks.Segmentation]:** It has the same functionality as the lightning task but it will also generate segmentation reports on demand.
 - **[`Classification Task`][quadra.tasks.Classification]:** This task is designed to train from scratch or finetune a classification model using the `pytorch-lightning` library.
 - **[`SklearnClassification Task`][quadra.tasks.classification.SklearnClassification]:** This task is designed to train an `sklearn` classifier on top of a torch feature extractor.
+- **[`PatchSklearnClassification Task`][quadra.tasks.patch.PatchSklearnClassification]:** This task is designed to train an `sklearn` patch classifier on top of a torch feature extractor.
 - **[`Anomalib Detection Task`][quadra.tasks.AnomalibDetection]:** This task is designed to train an anomaly detection model using the `anomalib` library.
 - **[`SSL (Self Supervised Learning) Task`][quadra.tasks.SSL]:** This task is designed to train a torch module with a given SSL algorithm.
 
 Most of these tasks have an associated evaluation task used for inference.
 ## Adding new tasks
 
-If your require the `pytorch-lightning` library, you can add a new task by extending the [`LightningTask`][quadra.base.LightningTask] class. Otherwise, you can simply start implementing a new task by extending the `Task` class.
\ No newline at end of file
+If you require the `pytorch-lightning` library, you can add a new task by extending the [`LightningTask`][quadra.tasks.base.LightningTask] class. Otherwise, you can simply start implementing a new task by extending the `Task` class.
\ No newline at end of file
diff --git a/docs_overrides/home.html b/docs_overrides/home.html
new file mode 100644
index 00000000..ff7b3fb2
--- /dev/null
+++ b/docs_overrides/home.html
@@ -0,0 +1,229 @@
+{% extends "main.html" %}
+
+{% block styles %}
+{{ super() }}
+<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/tw-elements/dist/css/tw-elements.min.css" />
+<!-- Additional styles for landing page -->
+<style>
+    [data-md-color-scheme="default"] {
+        --md-primary-bg-color: : #000;
+        --md-default-fg-color: #000;
+        --quadra-logo-single: #000;
+    }
+
+    [data-md-color-scheme="slate"] {
+        --quadra-logo-single: #FFF;
+    }
+
+    .q-logo-single {
+        fill: var(--quadra-logo-single);
+    }
+
+    /* Application header should be static for the landing page */
+    .md-header {
+        position: initial;
+        /* remove background color */
+        background-color: transparent;
+    }
+
+
+    .md-tabs {
+        background-color: transparent;
+    }
+
+    /* Remove spacing, as we cannot hide it completely */
+    .md-main__inner {
+        margin: 0;
+    }
+
+    /* Hide main content for now */
+    .md-content {
+        display: none;
+    }
+
+    /* Force header title visibility */
+    .md-header__title--active .md-header__topic {
+        opacity: 1;
+    }
+
+    .md-header__title--active .md-header__topic+.md-header__topic {
+        opacity: 0;
+    }
+
+    /* Hide table of contents */
+    @media screen and (min-width: 60em) {
+        .md-sidebar--secondary {
+            display: none;
+        }
+    }
+
+    @media screen and (min-width: 76.25em) {
+        .md-sidebar--primary {
+            display: none;
+        }
+    }
+
+    /* Enable navigation  */
+    @media screen and (min-width: 60em) {
+        .md-sidebar--primary {
+            display: block;
+        }
+    }
+
+    .lp-rolling-text-mask {
+        height: calc(1em + var(--rolling-text-offset));
+        position: relative;
+        overflow: hidden;
+        margin-top: var(--rolling-text-offset);
+        display: flex;
+        align-items: center;
+        justify-content: center;
+    }
+
+    .lp-rolling-text-mask span {
+        display: block;
+        box-sizing: border-box;
+        position: absolute;
+        top: 0;
+        transform: translateY(1.25em);
+        padding-bottom: var(--rolling-text-offset);
+        -webkit-background-clip: text;
+        background-clip: text;
+        -webkit-text-fill-color: transparent;
+        background-repeat: no-repeat;
+        white-space: nowrap;
+    }
+
+    .lp-rolling-text-mask span[data-show] {
+        transform: translateY(0);
+        transition: transform 0.8s ease-in-out;
+    }
+
+    .lp-rolling-text-mask span[data-up] {
+        transform: translateY(-1.25em);
+        transition: transform 0.8s ease-in-out;
+    }
+
+    .lp-rolling-text-mask span:nth-child(1) {
+        background-image: linear-gradient(45deg, var(--quadra-blue) 50%, var(--quadra-blue-dark));
+    }
+
+    .lp-rolling-text-mask span:nth-child(2) {
+        background-image: linear-gradient(45deg, var(--quadra-green) 50%, var(--quadra-green-dark));
+    }
+
+    .lp-rolling-text-mask span:nth-child(3) {
+        background-image: linear-gradient(45deg, var(--quadra-purple) 50%, var(--quadra-purple-dark));
+    }
+
+    .lp-rolling-text-mask span:nth-child(4) {
+        background-image: linear-gradient(45deg, var(--quadra-orange) 50%, var(--quadra-orange-dark));
+    }
+
+    .lp-hero-explain {
+        margin-top: -2em;
+        margin-left: 15%;
+        margin-right: 15%;
+        opacity: 0.7;
+    }
+
+    .lp-hero-button-container {
+        margin-top: 1.5em;
+        display: flex;
+        justify-content: center;
+    }
+
+    .lp-hero-button {
+        margin-top: 2%;
+        margin-left: 1%;
+        margin-right: 1%;
+        margin-bottom: 2%;
+    }
+
+    .lp-quadra-container {
+        justify-content: center;
+        height: 100vh;
+        height: 100vh;
+
+    }
+</style>
+{% endblock %}
+
+{% block tabs %}
+{{ super() }}
+
+<div class="lp-quadra-container py-24 px-6 text-center">
+    <h1 class="mt-2 mb-12 text-2xl font-bold tracking-tight md:text-5xl xl:text-6xl sm:text-2xl">
+
+        Effortless deep learning <br />
+        experiments with:
+        <div class="lp-rolling-text-mask">
+            <span data-show>Configuration files.</span>
+            <span>Reproducible workflows.</span>
+            <span>Model manager.</span>
+            <span>Modular tasks.</span>
+        </div>
+    </h1>
+    <p class="text-base lp-hero-explain sm:text-sm ">
+        Quadra simplifies handling many deep learning experiments by skipping boilerplate configurations.
+    </p>
+
+    <div class="md-hero__inner lp-hero_button-container md-typeset">
+        <a href="{{ page.next_page.url | url }}" title="{{ page.next_page.title | striptags }}"
+            class="md-button md-button--primary lp-hero-button">Learn More</a>
+        <a href="{{ config.repo_url }}" title="{{ lang.t('source.link.title') }}"
+            class="md-button md-button--secondary lp-hero-button">GitHub</a>
+    </div>
+
+    <hr
+        class="my-12 h-px border-t-0 bg-transparent bg-gradient-to-r from-transparent via-neutral-500 to-transparent opacity-25 dark:opacity-100" />
+
+</div>
+
+
+<script>
+    const ImgToSvg = async (img) => {
+        const s = document.createElement('div');
+        s.innerHTML = await (await fetch(img.src)).text();
+        // remove style tag from innerhtml
+        s.innerHTML = s.innerHTML.replace(/<style([\s\S]*?)<\/style>/gi, '');
+        s.firstChild.classList = img.classList;
+        img.replaceWith(s.firstChild);
+    };
+
+    window.addEventListener('DOMContentLoaded', async () => {
+        const images = document.querySelectorAll('img[alt="logo"]');
+        for (const img of images) {
+            await ImgToSvg(img);
+        }
+    });
+</script>
+
+<script src="https://cdn.jsdelivr.net/npm/tw-elements/dist/js/tw-elements.umd.min.js"></script>
+
+<script>
+
+    setInterval(function () {
+        const rollingTextMask = document.querySelector('div.lp-rolling-text-mask');
+        const show = rollingTextMask.querySelector('span[data-show]');
+        const next = show.nextElementSibling || rollingTextMask.querySelector('span:first-child');
+        const up = rollingTextMask.querySelector('span[data-up]');
+
+        if (up) {
+            up.removeAttribute('data-up');
+        }
+
+        show.removeAttribute('data-show');
+        show.setAttribute('data-up', '');
+        next.setAttribute('data-show', '');
+    }, 2000);
+
+</script>
+
+{% endblock %}
+
+
+<!-- Scripts -->
+{% block scripts %}
+{{ super() }}
+{% endblock %}
\ No newline at end of file
diff --git a/mkdocs.yml b/mkdocs.yml
index b2bff8e7..ae89f3d2 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -1,16 +1,20 @@
 site_name: Quadra
+site_url: https://orobix.github.io/quadra/
+site_description: Quadra
+repo_url: https://github.com/orobix/quadra
+repo_name: Orobix/quadra
 use_directory_urls: false
 nav:
-  - index.md
-  # defer to gen-files + literate-nav
+  - Home: index.md
+  - Getting Started: getting_started.md
   - Code Reference: reference/
   - Tutorials:
-      - Getting Started: tutorials/install.md
       - Datamodules: tutorials/datamodules.md
       - Tasks: tutorials/tasks.md
       - Configurations: tutorials/configurations.md
       - Setting up devices: tutorials/devices_setup.md
       - Reproducibility: tutorials/reproducibility.md
+      - Model Management: tutorials/model_management.md
       - External projects integration: tutorials/integration.md
       - Contributing: tutorials/contribution.md
       - Building documentation: tutorials/documentation.md
@@ -33,29 +37,32 @@ theme:
   name: "material"
   locale: en
   custom_dir: docs_overrides
-  logo: images/orobix_logo.svg
+  logo: images/quadra_logo.svg
+  icon:
+    repo: fontawesome/brands/github
   font:
     text: Roboto
     code: Roboto Mono
   palette:
-    - media: "(prefers-color-scheme: light)"
-      scheme: default
-      primary: deep orange
-      accent: orange
+    - scheme: default
+      primary: teal
+      accent: light blue
       toggle:
         icon: material/weather-night
         name: Switch to dark mode
-    - media: "(prefers-color-scheme: dark)"
-      scheme: slate
-      primary: deep orange
-      accent: orange
+
+    - scheme: slate
+      primary: teal
+      accent: light blue
       toggle:
         icon: material/weather-sunny
         name: Switch to light mode
+
   features:
     - content.code.annotate
     - content.code.copy
     - navigation.tabs
+    - navigation.tabs.sticky
     - navigation.top
 
 plugins:
@@ -92,9 +99,6 @@ extra:
       - latest
   generator: false
 
-extra_css:
-  - css/mkdocstrings.css
-
 copyright: Copyright &copy; 2023 Orobix
 
 markdown_extensions:
@@ -135,6 +139,11 @@ markdown_extensions:
       custom_checkbox: true
   - pymdownx.tilde
 
+extra_css:
+  - css/colors.css
+  - css/mkdocstrings.css
+  - css/landing.css
+
 extra_javascript:
   - javascripts/mathjax.js
   - https://polyfill.io/v3/polyfill.min.js?features=es6
diff --git a/pyproject.toml b/pyproject.toml
index 7cdb048b..bcff07ab 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "quadra"
-version = "1.0.2"
+version = "1.1.0"
 description = "Deep Learning experiment orchestration library"
 authors = [
   {name = "Alessandro Polidori", email = "alessandro.polidori@orobix.com"},
@@ -63,7 +63,8 @@ dependencies = [
   "h5py==3.8.*",
   "timm==0.6.12", # required by smp
   "segmentation-models-pytorch==0.3.*",
-  "anomalib@git+https://github.com/orobix/anomalib.git@v0.4.0+obx.1.0.1"
+  "anomalib@git+https://github.com/orobix/anomalib.git@v0.4.0+obx.1.0.1",
+  "xxhash==3.2.*",
 ]
 
 [project.optional-dependencies]
@@ -90,13 +91,14 @@ dev = [
 ]
 
 docs = [
-    "mkdocs==1.4.*",
-    "mkdocs-material==9.0.*",
-    "mkdocstrings-python==0.8.*",
-    "mkdocs-gen-files==0.4.*",
-    "mkdocs-literate-nav==0.6.*",
-    "mkdocs-section-index==0.3.*",
-    "mike==1.1.*",
+    "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"
 ]
 
 [tool.setuptools]
@@ -115,7 +117,7 @@ repository = "https://github.com/orobix/quadra"
 
 # Adapted from https://realpython.com/pypi-publish-python-package/#version-your-package
 [tool.bumpver]
-current_version = "1.0.2"
+current_version = "1.1.0"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit_message  = "build: Bump version {old_version} -> {new_version}"
 commit          = true
diff --git a/quadra/__init__.py b/quadra/__init__.py
index 21677a1b..0f5fca3e 100644
--- a/quadra/__init__.py
+++ b/quadra/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "1.0.2"
+__version__ = "1.1.0"
 
 
 def get_version():
diff --git a/quadra/callbacks/mlflow.py b/quadra/callbacks/mlflow.py
index c77f9adc..17b7b6f4 100644
--- a/quadra/callbacks/mlflow.py
+++ b/quadra/callbacks/mlflow.py
@@ -9,6 +9,8 @@
 from pytorch_lightning.utilities import rank_zero_only
 from pytorch_lightning.utilities.types import STEP_OUTPUT
 
+from quadra.utils.mlflow import get_mlflow_logger
+
 
 def check_minio_credentials() -> None:
     """Check minio credentials for aws based storage such as minio.
@@ -41,38 +43,28 @@ def check_file_server_dependencies() -> None:
         ) from e
 
 
-def set_minio_credentials(access_key_id: str, secret_access_key: str) -> None:
-    """Setting minio credential as environment variables.
+def validate_artifact_storage(logger: MLFlowLogger):
+    """Validate artifact storage.
 
     Args:
-        access_key_id: string id
-        secret_access_key: string key.
+        logger: Mlflow logger from pytorch lightning.
 
-    Returns:
-        None
     """
-    os.environ["AWS_ACCESS_KEY_ID"] = access_key_id
-    os.environ["AWS_SECRET_ACCESS_KEY"] = secret_access_key
-
+    from quadra.utils.utils import get_logger  # pylint: disable=[import-outside-toplevel]
 
-def get_mlflow_logger(trainer: Trainer) -> Optional[MLFlowLogger]:
-    """Safely get Mlflow logger from Trainer loggers.
+    log = get_logger(__name__)
 
-    Args:
-        trainer: Pytorch Lightning trainer.
-
-    Returns:
-        An mlflow logger if available, else None.
-    """
-    if isinstance(trainer.logger, MLFlowLogger):
-        return trainer.logger
-
-    if isinstance(trainer.logger, list):
-        for logger in trainer.logger:
-            if isinstance(logger, MLFlowLogger):
-                return logger
-
-    return None
+    client = logger.experiment
+    # TODO: we have to access the internal api to get the artifact uri, however there could be a better way
+    artifact_uri = client._tracking_client._get_artifact_repo(  # pylint: disable=protected-access
+        logger.run_id
+    ).artifact_uri
+    if artifact_uri.startswith("s3://"):
+        check_minio_credentials()
+        check_file_server_dependencies()
+        log.info("Mlflow artifact storage is AWS/S3 basedand credentials and dependencies are satisfied.")
+    else:
+        log.info("Mlflow artifact storage uri is %s. Validation checks are not implemented.", artifact_uri)
 
 
 class UploadCodeAsArtifact(Callback):
@@ -88,7 +80,6 @@ class UploadCodeAsArtifact(Callback):
 
     def __init__(self, source_dir: str):
         self.source_dir = source_dir
-        check_file_server_dependencies()
 
     @rank_zero_only
     def on_test_end(self, trainer: Trainer, pl_module: LightningModule):
@@ -220,7 +211,6 @@ def __init__(
         self.ckpt_ext = ckpt_ext
         self.delete_after_upload = delete_after_upload
         self.upload = upload
-        check_file_server_dependencies()
 
     @rank_zero_only
     def on_test_end(self, trainer: Trainer, pl_module: LightningModule):
diff --git a/quadra/configs/callbacks/all.yaml b/quadra/configs/callbacks/all.yaml
index 4a6fae16..f8e75071 100644
--- a/quadra/configs/callbacks/all.yaml
+++ b/quadra/configs/callbacks/all.yaml
@@ -11,7 +11,7 @@ model_checkpoint:
 
 upload_code_as_artifact:
   _target_: quadra.callbacks.mlflow.UploadCodeAsArtifact
-  source_dir: ${work_dir}/src/
+  source_dir: ${work_dir}/quadra/
 
 upload_ckpts_as_artifact:
   _target_: quadra.callbacks.mlflow.UploadCheckpointsAsArtifact
diff --git a/quadra/configs/experiment/base/anomaly/padim.yaml b/quadra/configs/experiment/base/anomaly/padim.yaml
index 596c8ca2..f9b1964b 100644
--- a/quadra/configs/experiment/base/anomaly/padim.yaml
+++ b/quadra/configs/experiment/base/anomaly/padim.yaml
@@ -21,8 +21,6 @@ datamodule:
   category:
   phase: train
 
-print_config: false
-
 core:
   tag: "run"
   name: padim_${model.model.backbone}_${trainer.max_epochs}
diff --git a/quadra/configs/experiment/base/anomaly/patchcore.yaml b/quadra/configs/experiment/base/anomaly/patchcore.yaml
index c754dd1e..1430f35c 100644
--- a/quadra/configs/experiment/base/anomaly/patchcore.yaml
+++ b/quadra/configs/experiment/base/anomaly/patchcore.yaml
@@ -21,8 +21,6 @@ datamodule:
   category:
   phase: train
 
-print_config: false
-
 core:
   tag: "run"
   name: patchcore_${model.model.backbone}_${trainer.max_epochs}
diff --git a/quadra/configs/experiment/base/classification/classification.yaml b/quadra/configs/experiment/base/classification/classification.yaml
index b35a8121..3ce38c9b 100644
--- a/quadra/configs/experiment/base/classification/classification.yaml
+++ b/quadra/configs/experiment/base/classification/classification.yaml
@@ -18,15 +18,17 @@ print_config: true
 
 model:
   num_classes: ???
-  gradcam: true
   module:
     lr_scheduler_interval: "epoch"
 
 task:
   lr_multiplier: 0.0
   run_test: True
-  export_type: [torchscript]
+  export_config:
+    types: [torchscript, pytorch]
+    input_shapes:
   report: True
+  gradcam: True
   output:
     example: True
 
diff --git a/quadra/configs/experiment/base/classification/classification_evaluation.yaml b/quadra/configs/experiment/base/classification/classification_evaluation.yaml
index a42060d9..439677ff 100644
--- a/quadra/configs/experiment/base/classification/classification_evaluation.yaml
+++ b/quadra/configs/experiment/base/classification/classification_evaluation.yaml
@@ -19,6 +19,7 @@ logger:
     run_name: ${core.name}
 
 task:
+  gradcam: true
   output:
     example: true
   model_path: ???
diff --git a/quadra/configs/experiment/base/classification/sklearn_classification.yaml b/quadra/configs/experiment/base/classification/sklearn_classification.yaml
index 222d16ce..1a833b99 100644
--- a/quadra/configs/experiment/base/classification/sklearn_classification.yaml
+++ b/quadra/configs/experiment/base/classification/sklearn_classification.yaml
@@ -13,6 +13,11 @@ backbone:
     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/base/segmentation/smp_evaluation.yaml b/quadra/configs/experiment/base/segmentation/smp_evaluation.yaml
new file mode 100644
index 00000000..37515ed4
--- /dev/null
+++ b/quadra/configs/experiment/base/segmentation/smp_evaluation.yaml
@@ -0,0 +1,17 @@
+# @package _global_
+
+defaults:
+  - override /datamodule: base/segmentation
+  - override /transforms: default_resize
+  - override /task: segmentation_evaluation
+core:
+  tag: "run"
+  name: "quadra_default"
+  upload_artifacts: True
+
+task:
+  model_path: ???
+
+datamodule:
+  num_workers: 5
+  batch_size: 32
diff --git a/quadra/configs/experiment/base/segmentation/smp_multiclass.yaml b/quadra/configs/experiment/base/segmentation/smp_multiclass.yaml
index 9611e41d..4c0d14bf 100644
--- a/quadra/configs/experiment/base/segmentation/smp_multiclass.yaml
+++ b/quadra/configs/experiment/base/segmentation/smp_multiclass.yaml
@@ -11,16 +11,16 @@ defaults:
   - override /backbone: smp
   - override /trainer: lightning_gpu
 
-core:
-  tag: "run"
-  name: "quadra_default"
-  upload_artifacts: True
-
 trainer:
   devices: [0]
   max_epochs: 100
   num_sanity_val_steps: 0
 
+core:
+  tag: "run"
+  name: "quadra_default"
+  upload_artifacts: True
+
 datamodule:
   num_workers: 8
   batch_size: 32
diff --git a/quadra/configs/experiment/base/segmentation/smp_multiclass_evaluation.yaml b/quadra/configs/experiment/base/segmentation/smp_multiclass_evaluation.yaml
new file mode 100644
index 00000000..5d249b7f
--- /dev/null
+++ b/quadra/configs/experiment/base/segmentation/smp_multiclass_evaluation.yaml
@@ -0,0 +1,18 @@
+# @package _global_
+
+defaults:
+  - override /datamodule: base/segmentation_multiclass
+  - override /transforms: default_resize
+  - override /task: segmentation_evaluation
+
+core:
+  tag: "run"
+  name: "quadra_default"
+  upload_artifacts: True
+
+task:
+  model_path: ???
+
+datamodule:
+  num_workers: 8
+  batch_size: 32
diff --git a/quadra/configs/logger/mlflow.yaml b/quadra/configs/logger/mlflow.yaml
index 812ff0b5..b59e22c9 100644
--- a/quadra/configs/logger/mlflow.yaml
+++ b/quadra/configs/logger/mlflow.yaml
@@ -9,3 +9,4 @@ mlflow:
   prefix: ""
   artifact_location: null
   run_id: null
+  log_model: false
diff --git a/quadra/configs/model/classification.yaml b/quadra/configs/model/classification.yaml
index 02e3d1c2..745806c5 100644
--- a/quadra/configs/model/classification.yaml
+++ b/quadra/configs/model/classification.yaml
@@ -8,5 +8,4 @@ classifier:
 module:
   _target_: quadra.modules.classification.ClassificationModule
   lr_scheduler_interval: "epoch"
-  criterion: ${loss}
-  gradcam: true
+  criterion: ${loss}
\ No newline at end of file
diff --git a/quadra/configs/task/anomalib/cfa.yaml b/quadra/configs/task/anomalib/cfa.yaml
index 8099f704..f675715c 100644
--- a/quadra/configs/task/anomalib/cfa.yaml
+++ b/quadra/configs/task/anomalib/cfa.yaml
@@ -1,6 +1,9 @@
 _target_: quadra.tasks.anomaly.AnomalibDetection
 report: true
 run_test: true
-export_type: [torchscript]
 module_function:
   _target_: anomalib.models.cfa.lightning_model.CfaLightning
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
+
diff --git a/quadra/configs/task/anomalib/cflow.yaml b/quadra/configs/task/anomalib/cflow.yaml
index ebf7d496..9cd0912b 100644
--- a/quadra/configs/task/anomalib/cflow.yaml
+++ b/quadra/configs/task/anomalib/cflow.yaml
@@ -1,6 +1,9 @@
 _target_: quadra.tasks.anomaly.AnomalibDetection
 report: true
 run_test: true
-export_type: null
 module_function:
   _target_: anomalib.models.cflow.lightning_model.CflowLightning
+export_config:
+  types:
+    - torchscript
+  input_shapes: null
diff --git a/quadra/configs/task/anomalib/csflow.yaml b/quadra/configs/task/anomalib/csflow.yaml
index e6219977..68a1ec24 100644
--- a/quadra/configs/task/anomalib/csflow.yaml
+++ b/quadra/configs/task/anomalib/csflow.yaml
@@ -1,6 +1,9 @@
 _target_: quadra.tasks.anomaly.AnomalibDetection
 report: true
 run_test: true
-export_type: [torchscript]
 module_function:
   _target_: anomalib.models.csflow.lightning_model.CsflowLightning
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
+
diff --git a/quadra/configs/task/anomalib/draem.yaml b/quadra/configs/task/anomalib/draem.yaml
index 7783b6fe..09d29e51 100644
--- a/quadra/configs/task/anomalib/draem.yaml
+++ b/quadra/configs/task/anomalib/draem.yaml
@@ -1,6 +1,9 @@
 _target_: quadra.tasks.anomaly.AnomalibDetection
 report: true
 run_test: true
-export_type: [torchscript]
 module_function:
   _target_: anomalib.models.draem.lightning_model.DraemLightning
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
+
diff --git a/quadra/configs/task/anomalib/fastflow.yaml b/quadra/configs/task/anomalib/fastflow.yaml
index ada2599c..d681d71d 100644
--- a/quadra/configs/task/anomalib/fastflow.yaml
+++ b/quadra/configs/task/anomalib/fastflow.yaml
@@ -1,6 +1,9 @@
 _target_: quadra.tasks.anomaly.AnomalibDetection
 report: true
 run_test: true
-export_type: [torchscript]
 module_function:
   _target_: anomalib.models.fastflow.lightning_model.FastflowLightning
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
+
diff --git a/quadra/configs/task/anomalib/padim.yaml b/quadra/configs/task/anomalib/padim.yaml
index 929eca85..25deb46d 100644
--- a/quadra/configs/task/anomalib/padim.yaml
+++ b/quadra/configs/task/anomalib/padim.yaml
@@ -1,6 +1,9 @@
 _target_: quadra.tasks.anomaly.AnomalibDetection
 report: true
 run_test: true
-export_type: [torchscript]
 module_function:
   _target_: anomalib.models.padim.lightning_model.PadimLightning
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
+
diff --git a/quadra/configs/task/anomalib/patchcore.yaml b/quadra/configs/task/anomalib/patchcore.yaml
index e6945c44..4dfdc376 100644
--- a/quadra/configs/task/anomalib/patchcore.yaml
+++ b/quadra/configs/task/anomalib/patchcore.yaml
@@ -1,6 +1,9 @@
 _target_: quadra.tasks.anomaly.AnomalibDetection
 report: true
 run_test: true
-export_type: [torchscript]
 module_function:
   _target_: anomalib.models.patchcore.lightning_model.PatchcoreLightning
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
+
diff --git a/quadra/configs/task/classification.yaml b/quadra/configs/task/classification.yaml
index 817bbe73..a2ecb8e3 100644
--- a/quadra/configs/task/classification.yaml
+++ b/quadra/configs/task/classification.yaml
@@ -1,7 +1,9 @@
 _target_: quadra.tasks.Classification
-export_type: [torchscript]
 lr_multiplier: null
 output:
   example: false
 report: false
 run_test: true
+export_config:
+  types: [pytorch, torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
diff --git a/quadra/configs/task/classification_evaluation.yaml b/quadra/configs/task/classification_evaluation.yaml
index cd93d489..f04583d8 100644
--- a/quadra/configs/task/classification_evaluation.yaml
+++ b/quadra/configs/task/classification_evaluation.yaml
@@ -3,3 +3,4 @@ device: cuda:0
 output:
   example: true
 model_path: ???
+report: true
diff --git a/quadra/configs/task/segmentation.yaml b/quadra/configs/task/segmentation.yaml
index a15bc98e..51b0b2e3 100644
--- a/quadra/configs/task/segmentation.yaml
+++ b/quadra/configs/task/segmentation.yaml
@@ -1,10 +1,12 @@
 _target_: quadra.tasks.Segmentation
 num_viz_samples: 10
 checkpoint_path: null
-export_type: [torchscript]
 run_test: true
 report: false
 evaluate:
   analysis: false
   control: false
   corruption: false
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
diff --git a/quadra/configs/task/segmentation_evaluation.yaml b/quadra/configs/task/segmentation_evaluation.yaml
new file mode 100644
index 00000000..c11f44ea
--- /dev/null
+++ b/quadra/configs/task/segmentation_evaluation.yaml
@@ -0,0 +1,3 @@
+_target_: quadra.tasks.SegmentationAnalysisEvaluation
+device: cuda:0
+model_path: ???
diff --git a/quadra/configs/task/sklearn_classification.yaml b/quadra/configs/task/sklearn_classification.yaml
index c484c34f..feae7377 100644
--- a/quadra/configs/task/sklearn_classification.yaml
+++ b/quadra/configs/task/sklearn_classification.yaml
@@ -1,9 +1,11 @@
 _target_: quadra.tasks.SklearnClassification
 device: "cuda:0"
-export_type: [torchscript]
 output:
   folder: "classification_experiment"
   save_backbone: false
   report: true
   example: true
   test_full_data: true
+export_config:
+  types: [pytorch, torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
diff --git a/quadra/configs/task/sklearn_classification_patch.yaml b/quadra/configs/task/sklearn_classification_patch.yaml
index b38a1d7a..2a25c366 100644
--- a/quadra/configs/task/sklearn_classification_patch.yaml
+++ b/quadra/configs/task/sklearn_classification_patch.yaml
@@ -1,9 +1,11 @@
 _target_: quadra.tasks.PatchSklearnClassification
 device: cuda:2
-export_type: [torchscript]
 output:
   folder: classification_patch_experiment
   save_backbone: false
   report: true
   example: true
   reconstruction_method: major_voting
+export_config:
+  types: [torchscript]
+  input_shapes: # Redefine the input shape if not automatically inferred
diff --git a/quadra/configs/task/ssl.yaml b/quadra/configs/task/ssl.yaml
index cc32f5db..24ae88e4 100644
--- a/quadra/configs/task/ssl.yaml
+++ b/quadra/configs/task/ssl.yaml
@@ -1,3 +1,5 @@
 _target_: quadra.tasks.SSL
-export_type: [torchscript]
 run_test: false
+export_config:
+  types: [torchscript]
+  input_shapes:
diff --git a/quadra/datamodules/base.py b/quadra/datamodules/base.py
index 7b6796be..ac284d2e 100644
--- a/quadra/datamodules/base.py
+++ b/quadra/datamodules/base.py
@@ -1,18 +1,22 @@
+import multiprocessing as mp
+import multiprocessing.pool as mpp
 import os
 import pickle as pkl
+import typing
 from functools import wraps
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
+from typing import Any, Callable, Dict, Iterable, List, Literal, Optional, Sequence, Tuple, Union, cast
 
 import albumentations
 import numpy as np
 import pandas as pd
 import torch
+import xxhash
 from pytorch_lightning import LightningDataModule
+from tqdm import tqdm
 
 from quadra.utils import utils
 
 log = utils.get_logger(__name__)
-IMAGE_EXTENSIONS: List[str] = [".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".pbm", ".pgm", ".ppm", ".pxm", ".pnm"]
 TrainDataset = Union[torch.utils.data.Dataset, Sequence[torch.utils.data.Dataset]]
 ValDataset = Union[torch.utils.data.Dataset, Sequence[torch.utils.data.Dataset]]
 TestDataset = torch.utils.data.Dataset
@@ -46,6 +50,73 @@ def __new__(cls, name, bases, dct):
         return super().__new__(cls, name, bases, dct)
 
 
+def compute_file_content_hash(path: str, hash_size: Literal[32, 64, 128] = 64) -> str:
+    """Get hash of a file based on its content.
+
+    Args:
+        path: Path to the file.
+        hash_size: Size of the hash. Must be one of [32, 64, 128].
+
+    Returns:
+        The hash of the file.
+    """
+    with open(path, "rb") as f:
+        data = f.read()
+
+        if hash_size == 32:
+            file_hash = xxhash.xxh32(data, seed=42).hexdigest()
+        elif hash_size == 64:
+            file_hash = xxhash.xxh64(data, seed=42).hexdigest()
+        elif hash_size == 128:
+            file_hash = xxhash.xxh128(data, seed=42).hexdigest()
+        else:
+            raise ValueError(f"Invalid hash size {hash_size}. Must be one of [32, 64, 128].")
+
+    return file_hash
+
+
+def compute_file_size_hash(path: str, hash_size: Literal[32, 64, 128] = 64) -> str:
+    """Get hash of a file based on its size.
+
+    Args:
+        path: Path to the file.
+        hash_size: Size of the hash. Must be one of [32, 64, 128].
+
+    Returns:
+        The hash of the file.
+    """
+    data = str(os.path.getsize(path))
+
+    if hash_size == 32:
+        file_hash = xxhash.xxh32(data, seed=42).hexdigest()
+    elif hash_size == 64:
+        file_hash = xxhash.xxh64(data, seed=42).hexdigest()
+    elif hash_size == 128:
+        file_hash = xxhash.xxh128(data, seed=42).hexdigest()
+    else:
+        raise ValueError(f"Invalid hash size {hash_size}. Must be one of [32, 64, 128].")
+
+    return file_hash
+
+
+@typing.no_type_check
+def istarmap(self, func: Callable, iterable: Iterable, chunksize: int = 1):
+    # pylint: disable=all
+    """Starmap-version of imap."""
+    self._check_running()
+    if chunksize < 1:
+        raise ValueError("Chunksize must be 1+, not {0:n}".format(chunksize))
+
+    task_batches = mpp.Pool._get_tasks(func, iterable, chunksize)
+    result = mpp.IMapIterator(self)
+    self._taskqueue.put((self._guarded_task_generation(result._job, mpp.starmapstar, task_batches), result._set_length))
+    return (item for chunk in result for item in chunk)
+
+
+# Patch Pool class to include istarmap
+mpp.Pool.istarmap = istarmap  # type: ignore[attr-defined]
+
+
 class BaseDataModule(LightningDataModule, metaclass=DecorateParentMethod):
     """Base class for all data modules.
 
@@ -61,6 +132,10 @@ class BaseDataModule(LightningDataModule, metaclass=DecorateParentMethod):
             Defaults to None.
         test_transform: Transformations for test dataset.
             Defaults to None.
+        enable_hashing: Whether to enable hashing of images. Defaults to True.
+        hash_size: Size of the hash. Must be one of [32, 64, 128]. Defaults to 64.
+        hash_type: Type of hash to use, if content hash is used, the hash is computed on the file content, otherwise
+            the hash is computed on the file size which is faster but less safe. Defaults to "content".
     """
 
     def __init__(
@@ -78,6 +153,9 @@ def __init__(
         train_transform: Optional[albumentations.Compose] = None,
         val_transform: Optional[albumentations.Compose] = None,
         test_transform: Optional[albumentations.Compose] = None,
+        enable_hashing: bool = True,
+        hash_size: Literal[32, 64, 128] = 64,
+        hash_type: Literal["content", "size"] = "content",
     ):
         super().__init__()
         self.num_workers = num_workers
@@ -88,6 +166,12 @@ def __init__(
         self.train_transform = train_transform
         self.val_transform = val_transform
         self.test_transform = test_transform
+        self.enable_hashing = enable_hashing
+        self.hash_size = hash_size
+        self.hash_type = hash_type
+
+        if self.hash_size not in [32, 64, 128]:
+            raise ValueError(f"Invalid hash size {self.hash_size}. Must be one of [32, 64, 128].")
 
         self.load_aug_images = load_aug_images
         self.aug_name = aug_name
@@ -102,6 +186,7 @@ def __init__(
         self.data_folder = "data"
         os.makedirs(self.data_folder, exist_ok=True)
         self.datamodule_checkpoint_file = os.path.join(self.data_folder, "datamodule.pkl")
+        self.dataset_file = os.path.join(self.data_folder, "dataset.csv")
 
     @property
     def train_data(self) -> pd.DataFrame:
@@ -166,12 +251,35 @@ def _prepare_data(self) -> None:
             "contained in the prepare_data method of a LightningModule."
         )
 
+    def hash_data(self) -> None:
+        """Computes the hash of the files inside the datasets."""
+        if not self.enable_hashing:
+            return
+
+        # TODO: We need to find a way to annotate the columns of data.
+        paths_and_hash_length = zip(self.data["samples"], [self.hash_size] * len(self.data))
+
+        with mp.Pool(min(8, mp.cpu_count() - 1)) as pool:
+            self.data["hash"] = list(
+                tqdm(
+                    pool.istarmap(  # type: ignore[attr-defined]
+                        compute_file_content_hash if self.hash_type == "content" else compute_file_size_hash,
+                        paths_and_hash_length,
+                    ),
+                    total=len(self.data),
+                    desc="Computing hashes",
+                )
+            )
+
+        self.data["hash_type"] = self.hash_type
+
     def prepare_data(self) -> None:
         """Prepares the data, should be overridden by subclasses."""
         if hasattr(self, "data"):
             return
 
         self._prepare_data()
+        self.hash_data()
         self.save_checkpoint()
 
     def __getstate__(self) -> Dict[str, Any]:
@@ -190,13 +298,11 @@ def save_checkpoint(self) -> None:
         datamodule to disk because we can't assign attributes to the datamodule in prepare_data when working with
         multiple gpus.
         """
-        saved_to_disk = False
-        if not os.path.exists(self.datamodule_checkpoint_file):
+        if not os.path.exists(self.datamodule_checkpoint_file) and not os.path.exists(self.dataset_file):
             with open(self.datamodule_checkpoint_file, "wb") as f:
                 pkl.dump(self, f)
-            saved_to_disk = True
 
-        if saved_to_disk:
+            self.data.to_csv(self.dataset_file, index=False)
             log.info("Datamodule checkpoint saved to disk.")
 
         if "targets" in self.data and not isinstance(self.data["targets"].iloc[0], np.ndarray):
diff --git a/quadra/datamodules/classification.py b/quadra/datamodules/classification.py
index 399cc903..6dfcebc4 100644
--- a/quadra/datamodules/classification.py
+++ b/quadra/datamodules/classification.py
@@ -12,7 +12,7 @@
 from timm.data.parsers.parser_image_folder import find_images_and_targets
 from torch.utils.data import DataLoader
 
-from quadra.datamodules.base import IMAGE_EXTENSIONS, BaseDataModule
+from quadra.datamodules.base import BaseDataModule
 from quadra.datasets import ImageClassificationListDataset
 from quadra.datasets.classification import MultilabelClassificationDataset
 from quadra.utils import utils
@@ -135,7 +135,7 @@ def _find_images_and_targets(
     ) -> Tuple[List[Tuple[str, int]], Dict[str, int]]:
         """Collects the samples from item folders."""
         images_and_targets, class_to_idx = find_images_and_targets(
-            folder=root_folder, types=IMAGE_EXTENSIONS, class_to_idx=class_to_idx
+            folder=root_folder, types=utils.IMAGE_EXTENSIONS, class_to_idx=class_to_idx
         )
         return images_and_targets, class_to_idx
 
@@ -397,6 +397,7 @@ class SklearnClassificationDataModule(BaseDataModule):
         label_map: Dictionary of conversion btw folder name and label.
         train_split_file: Optional path to a csv file containing the train split samples.
         test_split_file: Optional path to a csv file containing the test split samples.
+        **kwargs: Additional arguments for BaseDataModule
     """
 
     def __init__(
@@ -422,6 +423,7 @@ def __init__(
         test_split_file: Optional[str] = None,
         name: str = "sklearn_classification_datamodule",
         dataset: Type[ImageClassificationListDataset] = ImageClassificationListDataset,
+        **kwargs: Any,
     ):
         super().__init__(
             data_path=data_path,
@@ -432,6 +434,7 @@ def __init__(
             train_transform=train_transform,
             val_transform=val_transform,
             test_transform=test_transform,
+            **kwargs,
         )
 
         self.class_to_idx = class_to_idx
diff --git a/quadra/datamodules/generic/oxford_pet.py b/quadra/datamodules/generic/oxford_pet.py
index fbd66717..b52342a7 100644
--- a/quadra/datamodules/generic/oxford_pet.py
+++ b/quadra/datamodules/generic/oxford_pet.py
@@ -1,5 +1,5 @@
 import os
-from typing import Optional, Type
+from typing import Any, Optional, Type
 
 import albumentations
 import cv2
@@ -44,6 +44,7 @@ def __init__(
         train_transform: Optional[albumentations.Compose] = None,
         test_transform: Optional[albumentations.Compose] = None,
         val_transform: Optional[albumentations.Compose] = None,
+        **kwargs: Any,
     ):
         super().__init__(
             data_path=data_path,
@@ -57,6 +58,7 @@ def __init__(
             val_transform=val_transform,
             batch_size=batch_size,
             num_workers=num_workers,
+            **kwargs,
         )
 
     _RESOURCES = (
diff --git a/quadra/datamodules/segmentation.py b/quadra/datamodules/segmentation.py
index a9ae46fe..5fc5058a 100644
--- a/quadra/datamodules/segmentation.py
+++ b/quadra/datamodules/segmentation.py
@@ -2,7 +2,7 @@
 import glob
 import os
 import random
-from typing import Dict, List, Optional, Tuple, Type
+from typing import Any, Dict, List, Optional, Tuple, Type
 
 import albumentations
 import cv2
@@ -59,6 +59,7 @@ def __init__(
         val_split_file: Optional[str] = None,
         num_data_class: Optional[int] = None,
         exclude_good: bool = False,
+        **kwargs: Any,
     ):
         super().__init__(
             data_path=data_path,
@@ -69,6 +70,7 @@ def __init__(
             train_transform=train_transform,
             val_transform=val_transform,
             test_transform=test_transform,
+            **kwargs,
         )
         self.test_size = test_size
         self.val_size = val_size
@@ -413,6 +415,7 @@ def __init__(
         exclude_good: bool = False,
         num_data_train: Optional[int] = None,
         one_hot_encoding: bool = False,
+        **kwargs: Any,
     ):
         super().__init__(
             data_path=data_path,
@@ -423,6 +426,7 @@ def __init__(
             train_transform=train_transform,
             val_transform=val_transform,
             test_transform=test_transform,
+            **kwargs,
         )
         self.test_size = test_size
         self.val_size = val_size
diff --git a/quadra/datasets/anomaly.py b/quadra/datasets/anomaly.py
index ce02f10c..b1058fcd 100644
--- a/quadra/datasets/anomaly.py
+++ b/quadra/datasets/anomaly.py
@@ -11,7 +11,7 @@
 from torch import Tensor
 from torch.utils.data import Dataset
 
-from quadra.datamodules.base import IMAGE_EXTENSIONS
+from quadra.utils.utils import IMAGE_EXTENSIONS
 
 
 def create_validation_set_from_test_set(samples: DataFrame, seed: int = 0) -> DataFrame:
diff --git a/quadra/datasets/patch.py b/quadra/datasets/patch.py
index 001c0497..f4d34c9e 100644
--- a/quadra/datasets/patch.py
+++ b/quadra/datasets/patch.py
@@ -91,7 +91,7 @@ def __init__(
     def __getitem__(self, idx) -> Tuple[np.ndarray, np.ndarray]:
         path, y = self.samples[idx]
 
-        h5_file = h5py.File(os.path.join(self.data_path, "train", path))
+        h5_file = h5py.File(path)
         x = cv2.imread(os.path.join(self.data_path, h5_file["img_path"][()].decode("utf-8")))
 
         weights = h5_file["triangles_weights"][()]
diff --git a/quadra/models/base.py b/quadra/models/base.py
new file mode 100644
index 00000000..a733e881
--- /dev/null
+++ b/quadra/models/base.py
@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+import inspect
+from typing import Any, Sequence
+
+import torch
+from torch import nn
+
+
+class ModelSignatureWrapper(nn.Module):
+    """Model wrapper used to retrieve input shape. It can be used as a decorator of nn.Module, the first call to the
+    forward method will retrieve the input shape and store it in the input_shapes attribute.
+    """
+
+    def __init__(self, model: nn.Module):
+        super().__init__()
+        self.instance = model
+        self.input_shapes: Any = None
+        self.disable = False
+
+        if isinstance(self.instance, ModelSignatureWrapper):
+            # Handle nested ModelSignatureWrapper
+            self.input_shapes = self.instance.input_shapes
+            self.instance = self.instance.instance
+
+    def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        """Retrieve the input shape and forward the model."""
+        if self.input_shapes is None and not self.disable:
+            try:
+                self.input_shapes = self._get_input_shapes(*args, **kwargs)
+            except Exception:
+                # Avoid circular import
+                # pylint: disable=import-outside-toplevel
+                from quadra.utils.utils import get_logger  # noqa
+
+                log = get_logger(__name__)
+                log.warning(
+                    "Failed to retrieve input shapes after forward! To export the model you'll need to "
+                    "provide the input shapes manually setting the export_config.input_shapes parameter! "
+                    "Alternatively you could try to use a forward with supported input types (and their compositions) "
+                    "(list, tuple, dict, tensors)."
+                )
+                self.disable = True
+
+        return self.instance.forward(*args, **kwargs)
+
+    def to(self, *args, **kwargs):
+        """Handle calls to to method returning the underlying model."""
+        return ModelSignatureWrapper(self.instance.to(*args, **kwargs))
+
+    def _get_input_shapes(self, *args: Any, **kwargs: Any) -> list[Any]:
+        """Retrieve the input shapes from the input. Inputs will be in the same order as the forward method
+        signature.
+        """
+        input_shapes = []
+
+        for arg in args:
+            input_shapes.append(self._get_input_shape(arg))
+
+        signature = inspect.signature(self.instance.forward)
+
+        for i, key in enumerate(signature.parameters.keys()):
+            if i < len(args):
+                continue
+
+            if key in kwargs:
+                input_shapes.append(self._get_input_shape(kwargs[key]))
+            else:
+                # Retrieve the default value
+                input_shapes.append(self._get_input_shape(signature.parameters[key].default))
+
+        return input_shapes
+
+    def _get_input_shape(self, inp: Sequence | torch.Tensor) -> list[Any] | tuple[Any, ...] | dict[str, Any]:
+        """Recursive function to retrieve the input shapes."""
+        if isinstance(inp, list):
+            return [self._get_input_shape(i) for i in inp]
+
+        if isinstance(inp, tuple):
+            return tuple(self._get_input_shape(i) for i in inp)
+
+        if isinstance(inp, torch.Tensor):
+            return tuple(inp.shape[1:])
+
+        if isinstance(inp, dict):
+            return {k: self._get_input_shape(v) for k, v in inp.items()}
+
+        raise ValueError(f"Input type {type(inp)} not supported")
+
+    def __getattr__(self, name: str) -> torch.Tensor | nn.Module:
+        if name in ["instance", "input_shapes"]:
+            return self.__dict__[name]
+
+        return getattr(self.__dict__["instance"], name)
+
+    def __setattr__(self, name: str, value: torch.Tensor | nn.Module) -> None:
+        if name in ["instance", "input_shapes"]:
+            self.__dict__[name] = value
+        else:
+            setattr(self.instance, name, value)
+
+    def __getattribute__(self, __name: str) -> Any:
+        if __name in ["instance", "input_shapes", "__dict__", "forward", "_get_input_shapes", "_get_input_shape", "to"]:
+            return super().__getattribute__(__name)
+
+        return getattr(self.instance, __name)
diff --git a/quadra/models/classification/backbones.py b/quadra/models/classification/backbones.py
index e26d7c15..2c5aca78 100644
--- a/quadra/models/classification/backbones.py
+++ b/quadra/models/classification/backbones.py
@@ -1,3 +1,5 @@
+from typing import Any
+
 import timm
 import torch
 from torch import nn
@@ -31,7 +33,7 @@ def __init__(
         freeze: bool = True,
         hyperspherical: bool = False,
         flatten_features: bool = True,
-        **torch_hub_kwargs: str
+        **torch_hub_kwargs: Any,
     ):
         self.pretrained = pretrained
         features_extractor = torch.hub.load(
@@ -70,7 +72,7 @@ def __init__(
         freeze: bool = True,
         hyperspherical: bool = False,
         flatten_features: bool = True,
-        **torchvision_kwargs: str
+        **torchvision_kwargs: Any,
     ):
         self.pretrained = pretrained
         model_function = models.__dict__[model_name]
@@ -110,7 +112,7 @@ def __init__(
         freeze: bool = True,
         hyperspherical: bool = False,
         flatten_features: bool = True,
-        **timm_kwargs: str
+        **timm_kwargs: Any,
     ):
         self.pretrained = pretrained
         features_extractor = timm.create_model(model_name, pretrained=self.pretrained, num_classes=0, **timm_kwargs)
diff --git a/quadra/modules/base.py b/quadra/modules/base.py
index 1f7ae004..ae030d26 100644
--- a/quadra/modules/base.py
+++ b/quadra/modules/base.py
@@ -8,6 +8,8 @@
 from torch import nn
 from torch.optim import Optimizer
 
+from quadra.models.base import ModelSignatureWrapper
+
 __all__ = ["BaseLightningModule", "SSLModule"]
 
 
@@ -28,7 +30,7 @@ def __init__(
         lr_scheduler_interval: Optional[str] = "epoch",
     ):
         super().__init__()
-        self.model = model
+        self.model = ModelSignatureWrapper(model)
         self.optimizer = optimizer
         self.schedulers = lr_scheduler
         self.lr_scheduler_interval = lr_scheduler_interval
@@ -90,7 +92,6 @@ def __init__(
         lr_scheduler: Optional[object] = None,
         lr_scheduler_interval: Optional[str] = "epoch",
     ):
-
         super().__init__(model, optimizer, lr_scheduler, lr_scheduler_interval)
         self.criterion = criterion
         self.classifier_train_loader: Optional[torch.utils.data.DataLoader]
@@ -176,7 +177,6 @@ def __init__(
         optimizer: Optional[Optimizer] = None,
         lr_scheduler: Optional[object] = None,
     ):
-
         super().__init__(model, optimizer, lr_scheduler)
         self.loss_fun = loss_fun
 
diff --git a/quadra/modules/classification/base.py b/quadra/modules/classification/base.py
index 5f29a793..c68685bb 100644
--- a/quadra/modules/classification/base.py
+++ b/quadra/modules/classification/base.py
@@ -1,15 +1,19 @@
 from typing import Any, List, Optional, Tuple, Union, cast
 
+import numpy as np
 import timm
 import torch
 import torchmetrics
 import torchmetrics.functional as TMF
 from pytorch_grad_cam import GradCAM
+from scipy import ndimage
 from torch import nn, optim
 
 from quadra.models.classification import BaseNetworkBuilder
 from quadra.modules.base import BaseLightningModule
+from quadra.utils.models import is_vision_transformer
 from quadra.utils.utils import get_logger
+from quadra.utils.vit_explainability import VitAttentionGradRollout
 
 log = get_logger(__name__)
 
@@ -45,11 +49,14 @@ def __init__(
         self.train_acc = torchmetrics.Accuracy()
         self.val_acc = torchmetrics.Accuracy()
         self.test_acc = torchmetrics.Accuracy()
-        self.cam: GradCAM
+        self.cam: Optional[GradCAM] = None
+        self.grad_rollout: Optional[VitAttentionGradRollout] = None
 
-        if not isinstance(self.model.features_extractor, timm.models.resnet.ResNet):
+        if not isinstance(self.model.features_extractor, timm.models.resnet.ResNet) and not is_vision_transformer(
+            cast(BaseNetworkBuilder, self.model).features_extractor
+        ):
             log.warning(
-                "Backbone must be compatible with gradcam, at the moment only ResNets supported, disabling gradcam"
+                "Backbone not compatible with gradcam. Only timm ResNets, timm ViTs and TorchHub dinoViTs supported",
             )
             self.gradcam = False
 
@@ -124,25 +131,55 @@ def test_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int):
             prog_bar=False,
         )
 
-    def on_predict_start(self) -> None:
-        """If gradcam will be computed, saves all requires_grad values and set them to True before the predict."""
-        if self.gradcam:
+    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]
-            self.cam = GradCAM(model=self.model, target_layers=target_layers, use_cuda=torch.cuda.is_available())
+            self.cam = GradCAM(
+                model=self.model,
+                target_layers=target_layers,
+                use_cuda=torch.cuda.is_available(),
+            )
+            # Activating gradients
+            for p in self.model.features_extractor.layer4[-1].parameters():
+                p.requires_grad = True
+        elif is_vision_transformer(cast(BaseNetworkBuilder, self.model).features_extractor):
+            self.grad_rollout = VitAttentionGradRollout(self.model)
+        else:
+            log.warning("Gradcam not implemented for this backbone, it won't be computed")
+            self.original_requires_grads.clear()
+            self.gradcam = False
 
+    def on_predict_start(self) -> None:
+        """If gradcam, prepares gradcam and saves params requires_grad state."""
+        if self.gradcam:
+            # Saving params requires_grad state
             for p in self.model.parameters():
                 self.original_requires_grads.append(p.requires_grad)
-                p.requires_grad = True
+            self.prepare_gradcam()
 
         return super().on_predict_start()
 
     def on_predict_end(self) -> None:
         """If we computed gradcam, requires_grad values are reset to original value."""
         if self.gradcam:
+            # Get back to initial state
             for i, p in enumerate(self.model.parameters()):
                 p.requires_grad = self.original_requires_grads[i]
 
-            self.cam.activations_and_grads.release()
+            # We are using GradCAM package only for resnets at the moment
+            if isinstance(self.model.features_extractor, timm.models.resnet.ResNet) and self.cam is not None:
+                # Needed to solve jitting bug
+                self.cam.activations_and_grads.release()
+            elif (
+                is_vision_transformer(cast(BaseNetworkBuilder, self.model).features_extractor)
+                and self.grad_rollout is not None
+            ):
+                for handle in self.grad_rollout.f_hook_handles:
+                    handle.remove()
+                for handle in self.grad_rollout.b_hook_handles:
+                    handle.remove()
+
         return super().on_predict_end()
 
     def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
@@ -158,18 +195,25 @@ def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> A
                 grayscale_cam: gray scale gradcams
         """
         im, _ = batch
-        # inference_mode set to false because gradcam needs gradients
+        outputs = self(im)
+        probs = torch.softmax(outputs, dim=1)
+        predicted_classes = torch.max(probs, dim=1).indices.tolist()
         if self.gradcam:
+            # inference_mode set to false because gradcam needs gradients
             with torch.inference_mode(False):
                 im = im.clone()
-                outputs = self(im)
-                probs = torch.softmax(outputs, dim=1)
-                predicted_classes = torch.max(probs, dim=1).indices.tolist()
-                grayscale_cam = self.cam(input_tensor=im, targets=None)
+
+                if isinstance(self.model.features_extractor, timm.models.resnet.ResNet) and self.cam:
+                    grayscale_cam = self.cam(input_tensor=im, targets=None)
+                elif (
+                    is_vision_transformer(cast(BaseNetworkBuilder, self.model).features_extractor) and self.grad_rollout
+                ):
+                    grayscale_cam_low_res = self.grad_rollout(input_tensor=im, targets_list=predicted_classes)
+                    orig_shape = grayscale_cam_low_res.shape
+                    new_shape = (orig_shape[0], im.shape[2], im.shape[3])
+                    zoom_factors = tuple(np.array(new_shape) / np.array(orig_shape))
+                    grayscale_cam = ndimage.zoom(grayscale_cam_low_res, zoom_factors, order=1)
         else:
-            outputs = self(im)
-            probs = torch.softmax(outputs, dim=1)
-            predicted_classes = torch.max(probs, dim=1).indices.tolist()
             grayscale_cam = None
         return predicted_classes, grayscale_cam
 
diff --git a/quadra/tasks/anomaly.py b/quadra/tasks/anomaly.py
index 383c5e06..3f36958a 100644
--- a/quadra/tasks/anomaly.py
+++ b/quadra/tasks/anomaly.py
@@ -2,8 +2,7 @@
 import json
 import os
 from collections import Counter
-from pathlib import Path
-from typing import Any, Dict, Generic, List, Optional, TypeVar, Union, cast
+from typing import Dict, Generic, List, Optional, TypeVar, Union, cast
 
 import cv2
 import hydra
@@ -21,10 +20,11 @@
 
 from quadra.callbacks.mlflow import get_mlflow_logger
 from quadra.datamodules import AnomalyDataModule
-from quadra.tasks.base import LightningTask, Task
+from quadra.modules.base import ModelSignatureWrapper
+from quadra.tasks.base import Evaluation, LightningTask
 from quadra.utils import utils
 from quadra.utils.classification import get_results
-from quadra.utils.export import export_torchscript_model, import_deployment_model
+from quadra.utils.export import export_torchscript_model
 
 log = utils.get_logger(__name__)
 
@@ -41,7 +41,11 @@ class AnomalibDetection(Generic[AnomalyDataModuleT], LightningTask[AnomalyDataMo
             Defaults to None.
         run_test: Whether to run the test after training. Defaults to False.
         report: Whether to report the results. Defaults to False.
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
@@ -51,14 +55,14 @@ def __init__(
         checkpoint_path: Optional[str] = None,
         run_test: bool = True,
         report: bool = True,
-        export_type: Optional[List[str]] = None,
+        export_config: Optional[DictConfig] = None,
     ):
         super().__init__(
             config=config,
             checkpoint_path=checkpoint_path,
             run_test=run_test,
             report=report,
-            export_type=export_type,
+            export_config=export_config,
         )
         self._module: AnomalyModule
         self.module_function = module_function
@@ -104,10 +108,11 @@ def prepare(self) -> None:
         """Prepare the task."""
         super().prepare()
         self.module = self.config.model
+        self.module.model = ModelSignatureWrapper(self.module.model)
 
     def export(self) -> None:
         """Export model for production."""
-        if self.export_type is None or len(self.export_type) == 0:
+        if self.export_config is None or len(self.export_config.types) == 0:
             log.info("No export type specified skipping export")
             return
 
@@ -115,19 +120,29 @@ def export(self) -> None:
             log.warning("Skipping export since fast_dev_run is enabled")
             return
 
-        input_width = self.config.transforms.get("input_width")
-        input_height = self.config.transforms.get("input_height")
-
         model = self.module.model
 
-        half_precision = self.trainer.precision == 16
+        input_shapes = self.export_config.input_shapes
+
+        half_precision = int(self.trainer.precision) == 16
 
-        for export_type in self.export_type:
+        for export_type in self.export_config.types:
             if export_type == "torchscript":
-                export_torchscript_model(model, (1, 3, input_height, input_width), self.export_folder, half_precision)
+                out = export_torchscript_model(
+                    model=model,
+                    input_shapes=input_shapes,
+                    output_path=self.export_folder,
+                    half_precision=half_precision,
+                )
+
+                if out is None:
+                    log.warning("Skipping torchscript export since the model is not supported")
+                    continue
+
+                _, input_shapes = out
 
         model_json = {
-            "input_size": [self.config.transforms.input_width, self.config.transforms.input_height, 3],
+            "input_size": input_shapes,
             "classes": {0: "good", 1: "defect"},
             "mean": list(self.config.transforms.mean),
             "std": list(self.config.transforms.std),
@@ -287,7 +302,7 @@ def _upload_artifacts(self):
                     utils.upload_file_tensorboard(a, tensorboard_logger)
 
 
-class AnomalibEvaluation(Task[AnomalyDataModule]):
+class AnomalibEvaluation(Evaluation[AnomalyDataModule]):
     """Evaluation task for Anomalib.
 
     Args:
@@ -301,63 +316,19 @@ class AnomalibEvaluation(Task[AnomalyDataModule]):
     def __init__(
         self, config: DictConfig, model_path: str, use_training_threshold: bool = False, device: Optional[str] = None
     ):
-        super().__init__(config=config)
-        self.model_data: Dict[str, Any]
-        self._deployment_model: Any
-        self.deployment_model_type: str
-        self.model_path = model_path
-        self.output_path = ""
-
-        if device is None:
-            self.device = utils.get_device()
-        else:
-            self.device = device
+        super().__init__(config=config, model_path=model_path, device=device)
 
-        self.config = config
-        self.report_path = ""
-        self.metadata = {"report_files": []}
-        self.model_info_filename = "model.json"
         self.use_training_threshold = use_training_threshold
 
-    @property
-    def deployment_model(self):
-        """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)
-
     def prepare(self) -> None:
         """Prepare the evaluation."""
-        with open(os.path.join(Path(self.model_path).parent, self.model_info_filename)) as f:
-            self.model_data = json.load(f)
-
-        if not isinstance(self.model_data, dict):
-            raise ValueError("Model info file is not a valid json")
-
-        if self.model_data["input_size"][0] != self.config.transforms.input_height:
-            log.warning(
-                f"Input height of the model ({self.model_data['input_size'][0]}) is different from the one specified "
-                + f"in the config ({self.config.transforms.input_height}). Fixing the config."
-            )
-            self.config.transforms.input_height = self.model_data["input_size"][0]
-
-        if self.model_data["input_size"][1] != self.config.transforms.input_width:
-            log.warning(
-                f"Input width of the model ({self.model_data['input_size'][1]}) is different from the one specified "
-                + f"in the config ({self.config.transforms.input_width}). Fixing the config."
-            )
-            self.config.transforms.input_width = self.model_data["input_size"][1]
-
-        self.deployment_model = self.model_path
-
         super().prepare()
+        self.datamodule = self.config.datamodule
 
     def test(self) -> None:
         """Perform test."""
         log.info("Running test")
+        # prepare_data() must be explicitly called because there is no lightning training
         self.datamodule.prepare_data()
         self.datamodule.setup(stage="test")
         test_dataloader = self.datamodule.test_dataloader()
@@ -422,8 +393,8 @@ def generate_report(self) -> None:
         if len(self.report_path) > 0:
             os.makedirs(self.report_path, exist_ok=True)
 
-        os.makedirs(os.path.join(self.output_path, "predictions"), exist_ok=True)
-        os.makedirs(os.path.join(self.output_path, "heatmaps"), exist_ok=True)
+        os.makedirs(os.path.join(self.report_path, "predictions"), exist_ok=True)
+        os.makedirs(os.path.join(self.report_path, "heatmaps"), exist_ok=True)
 
         anomaly_scores = self.metadata["anomaly_scores"].cpu().numpy()
         good_scores = anomaly_scores[np.where(np.array(self.metadata["image_labels"]) == 0)]
@@ -517,7 +488,7 @@ def generate_report(self) -> None:
 
             output_mask = output_mask * 255
             output_mask = cv2.resize(output_mask, (img.shape[1], img.shape[0]))
-            cv2.imwrite(os.path.join(self.output_path, "predictions", output_mask_name), output_mask)
+            cv2.imwrite(os.path.join(self.report_path, "predictions", output_mask_name), output_mask)
 
             # Normalize the heatmaps based on the current min and max anomaly score, otherwise even on good images the
             # anomaly map looks like there are defects while it's not true
@@ -529,7 +500,7 @@ def generate_report(self) -> None:
             output_heatmap = anomaly_map_to_color_map(output_heatmap, normalize=False)
             output_heatmap = cv2.resize(output_heatmap, (img.shape[1], img.shape[0]))
             cv2.imwrite(
-                os.path.join(self.output_path, "heatmaps", output_mask_name),
+                os.path.join(self.report_path, "heatmaps", output_mask_name),
                 cv2.cvtColor(output_heatmap, cv2.COLOR_RGB2BGR),
             )
 
@@ -541,7 +512,7 @@ def generate_report(self) -> None:
             ],
         }
 
-        with open(os.path.join(self.output_path, "anomaly_test_output.json"), "w") as f:
+        with open(os.path.join(self.report_path, "anomaly_test_output.json"), "w") as f:
             json.dump(json_output, f)
 
     def execute(self) -> None:
diff --git a/quadra/tasks/base.py b/quadra/tasks/base.py
index 9bedccd5..3b0f1ed5 100644
--- a/quadra/tasks/base.py
+++ b/quadra/tasks/base.py
@@ -1,4 +1,6 @@
+import json
 import os
+from pathlib import Path
 from typing import Any, Dict, Generic, List, Optional, TypeVar, Union
 
 import hydra
@@ -6,12 +8,15 @@
 from hydra.core.hydra_config import HydraConfig
 from omegaconf import DictConfig, OmegaConf
 from pytorch_lightning import Callback, LightningModule, Trainer
-from pytorch_lightning.loggers import Logger
+from pytorch_lightning.loggers import Logger, MLFlowLogger
 from pytorch_lightning.utilities.device_parser import parse_gpu_ids
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from torch import nn
 from torch.jit._script import RecursiveScriptModule
+from torch.nn import Module
 
 from quadra import get_version
+from quadra.callbacks.mlflow import validate_artifact_storage
 from quadra.datamodules.base import BaseDataModule
 from quadra.utils import utils
 from quadra.utils.export import import_deployment_model
@@ -24,13 +29,18 @@ class Task(Generic[DataModuleT]):
     """Base Experiment Task.
 
     Args:
-        config: The experiment configuration
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        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_type: Optional[List[str]] = None):
+    def __init__(self, config: DictConfig, export_config: Optional[DictConfig] = None):
         self.config = config
-        self.export_type = export_type
+        self.export_config = export_config
+        self.export_folder: str = "deployment_model"
         self._datamodule: DataModuleT
         self.metadata: Dict[str, Any]
         self.save_config()
@@ -82,7 +92,7 @@ def execute(self) -> None:
         self.prepare()
         self.train()
         self.test()
-        if self.export_type is not None and len(self.export_type) > 0:
+        if self.export_config is not None and len(self.export_config.types) > 0:
             self.export()
         self.generate_report()
         self.finalize()
@@ -96,7 +106,11 @@ class LightningTask(Generic[DataModuleT], Task[DataModuleT]):
         checkpoint_path: The path to the checkpoint to load the model from. Defaults to None.
         run_test: Whether to run the test after training. Defaults to False.
         report: Whether to generate a report. Defaults to False.
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
@@ -105,9 +119,9 @@ def __init__(
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
         report: bool = False,
-        export_type: Optional[List[str]] = None,
+        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config, export_type=export_type)
+        super().__init__(config, export_config=export_config)
         self.config = config
         self.checkpoint_path = checkpoint_path
         self.run_test = run_test
@@ -122,11 +136,13 @@ def prepare(self) -> None:
         """Prepare the experiment."""
         super().prepare()
 
+        # First setup loggers since some callbacks might need logger setup correctly.
+        if "logger" in self.config:
+            self.logger = self.config.logger
+
         if "callbacks" in self.config:
             self.callbacks = self.config.callbacks
 
-        if "logger" in self.config:
-            self.logger = self.config.logger
         self.devices = self.config.trainer.devices
         self.trainer = self.config.trainer
 
@@ -199,6 +215,8 @@ def logger(self, logger_config) -> None:
             if "_target_" in lg_conf:
                 log.info("Instantiating logger <%s>", lg_conf["_target_"])
                 logger = hydra.utils.instantiate(lg_conf)
+                if isinstance(logger, MLFlowLogger):
+                    validate_artifact_storage(logger)
                 instantiated_loggers.append(logger)
 
         self._logger = instantiated_loggers
@@ -233,6 +251,7 @@ def train(self) -> None:
             model=self.module,
             trainer=self.trainer,
         )
+
         self.trainer.fit(model=self.module, datamodule=self.datamodule)
 
     def test(self) -> Any:
@@ -245,11 +264,12 @@ def finalize(self) -> None:
         super().finalize()
         utils.finish(
             config=self.config,
-            model=self.module,
+            module=self.module,
             datamodule=self.datamodule,
             trainer=self.trainer,
             callbacks=self.callbacks,
             logger=self.logger,
+            export_folder=self.export_folder,
         )
 
         if not self.config.trainer.get("fast_dev_run"):
@@ -273,7 +293,7 @@ def execute(self) -> None:
         self.train()
         if self.run_test:
             self.test()
-        if self.export_type is not None and len(self.export_type) > 0:
+        if self.export_config is not None and len(self.export_config.types) > 0:
             self.export()
         if self.report:
             self.generate_report()
@@ -290,46 +310,74 @@ def execute(self) -> None:
         log.info("If you are reading this, it means that library is installed correctly!")
 
 
-class Evaluation(Task):
+class Evaluation(Generic[DataModuleT], Task[DataModuleT]):
     """Base Evaluation Task with deployment models.
 
     Args:
         config: The experiment configuration
         model_path: The model path.
-        report_folder: The report folder. Defaults to None.
+        device: Device to use for evaluation. If None, the device is automatically determined.
 
-    Raises:
-        ValueError: If the experiment path is not provided
     """
 
     def __init__(
         self,
         config: DictConfig,
         model_path: str,
-        report_folder: Optional[str] = None,
+        device: Optional[str] = None,
     ):
         super().__init__(config=config)
+
+        if device is None:
+            self.device = utils.get_device()
+        else:
+            self.device = device
+
         self.config = config
-        self.metadata = {"report_files": []}
+        self.model_data: Dict[str, Any]
         self.model_path = model_path
-        self.device = utils.get_device()
-        self.report_folder = report_folder
-        self._deployment_model: RecursiveScriptModule
+        self._deployment_model: Union[RecursiveScriptModule, Module]
         self.deployment_model_type: str
-        if self.report_folder is None:
-            log.warning("Report folder is not provided, using default report folder")
-            self.report_folder = "report"
+        self.model_info_filename = "model.json"
+        self.report_path = ""
+        self.metadata = {"report_files": []}
 
     @property
-    def deployment_model(self) -> RecursiveScriptModule:
-        """RecursiveScriptModule: The deployment model."""
+    def deployment_model(self) -> Union[RecursiveScriptModule, nn.Module]:
+        """Deployment model."""
         return self._deployment_model
 
     @deployment_model.setter
-    def deployment_model(self, model: RecursiveScriptModule) -> None:
-        """RecursiveScriptModule: The deployment model."""
-        self._deployment_model = model
+    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)
 
     def prepare(self) -> None:
         """Prepare the evaluation."""
-        self.deployment_model, self.deployment_model_type = import_deployment_model(self.model_path, self.device)
+        with open(os.path.join(Path(self.model_path).parent, self.model_info_filename)) as f:
+            self.model_data = json.load(f)
+
+        if not isinstance(self.model_data, dict):
+            raise ValueError("Model info file is not a valid json")
+
+        for input_size in self.model_data["input_size"]:
+            if len(input_size) != 3:
+                continue
+
+            # Adjust the transform for 2D models (CxHxW)
+            # We assume that each input size has the same height and width
+            if input_size[1] != self.config.transforms.input_height:
+                log.warning(
+                    f"Input height of the model ({input_size[1]}) is different from the one specified "
+                    + f"in the config ({self.config.transforms.input_height}). Fixing the config."
+                )
+                self.config.transforms.input_height = input_size[1]
+
+            if input_size[2] != self.config.transforms.input_width:
+                log.warning(
+                    f"Input width of the model ({input_size[2]}) is different from the one specified "
+                    + f"in the config ({self.config.transforms.input_width}). Fixing the config."
+                )
+                self.config.transforms.input_width = input_size[2]
+
+        self.deployment_model = self.model_path  # type: ignore[assignment]
diff --git a/quadra/tasks/classification.py b/quadra/tasks/classification.py
index 298818d5..cf6d7144 100644
--- a/quadra/tasks/classification.py
+++ b/quadra/tasks/classification.py
@@ -10,11 +10,15 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+import timm
 import torch
 from joblib import dump, load
-from omegaconf import DictConfig
+from omegaconf import DictConfig, OmegaConf
+from pytorch_grad_cam import GradCAM
+from scipy import ndimage
 from sklearn.base import ClassifierMixin
 from sklearn.metrics import ConfusionMatrixDisplay
+from torch import nn
 from tqdm import tqdm
 
 from quadra.callbacks.mlflow import get_mlflow_logger
@@ -25,13 +29,16 @@
     SklearnClassificationDataModule,
 )
 from quadra.datasets.classification import ImageClassificationListDataset
+from quadra.models.base import ModelSignatureWrapper
+from quadra.models.classification import BaseNetworkBuilder
 from quadra.modules.classification import ClassificationModule
-from quadra.tasks.base import LightningTask, Task
+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_torchscript_model, import_deployment_model
-from quadra.utils.models import get_feature
+from quadra.utils.export import export_pytorch_model, export_torchscript_model, import_deployment_model
+from quadra.utils.models import get_feature, is_vision_transformer
+from quadra.utils.vit_explainability import VitAttentionGradRollout
 
 log = utils.get_logger(__name__)
 
@@ -52,7 +59,12 @@ class Classification(Generic[ClassificationDataModuleT], LightningTask[Classific
     Args:
         config: The experiment configuration
         output: The otuput configuration.
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
+        gradcam: Whether to compute gradcams
         checkpoint_path: The path to the checkpoint to load the model from. Defaults to None.
         lr_multiplier: The multiplier for the backbone learning rate. Defaults to None.
         output: The ouput configuration (under task config). It contains the bool "example" to generate
@@ -67,21 +79,29 @@ def __init__(
         output: DictConfig,
         checkpoint_path: Optional[str] = None,
         lr_multiplier: Optional[float] = None,
-        export_type: Optional[List[str]] = None,
+        export_config: Optional[DictConfig] = None,
+        gradcam: bool = False,
         report: bool = False,
         run_test: bool = False,
     ):
         super().__init__(
-            config=config, checkpoint_path=checkpoint_path, run_test=run_test, report=report, export_type=export_type
+            config=config,
+            checkpoint_path=checkpoint_path,
+            run_test=run_test,
+            report=report,
+            export_config=export_config,
         )
         self.output = output
+        self.gradcam = gradcam
         self._lr_multiplier = lr_multiplier
-        self._pre_classifier: torch.nn.Module
-        self._classifier: torch.nn.Module
-        self._model: torch.nn.Module
+        self._pre_classifier: nn.Module
+        self._classifier: nn.Module
+        self._model: nn.Module
         self._optimizer: torch.optim.Optimizer
         self._scheduler: torch.optim.lr_scheduler._LRScheduler
+        self.model_json: Optional[Dict[str, Any]] = None
         self.export_folder: str = "deployment_model"
+        self.deploy_info_file: str = "model.json"
         self.report_confmat: pd.DataFrame
 
     @property
@@ -93,9 +113,9 @@ def optimizer(self) -> torch.optim.Optimizer:
     def optimizer(self, optimizer_config: DictConfig) -> None:
         """Set the optimizer."""
         if (
-            isinstance(self.model.features_extractor, torch.nn.Module)
-            and isinstance(self.model.pre_classifier, torch.nn.Module)
-            and isinstance(self.model.classifier, torch.nn.Module)
+            isinstance(self.model.features_extractor, nn.Module)
+            and isinstance(self.model.pre_classifier, nn.Module)
+            and isinstance(self.model.classifier, nn.Module)
         ):
             log.info("Instantiating optimizer <%s>", self.config.optimizer["_target_"])
             if self._lr_multiplier is not None and self._lr_multiplier > 0:
@@ -155,6 +175,7 @@ def module(self, module_config):
             model=self.model,
             optimizer=self.optimizer,
             lr_scheduler=self.scheduler,
+            gradcam=self.gradcam,
         )
         if self.checkpoint_path is not None:
             log.info("Loading model from lightning checkpoint: %s", self.checkpoint_path)
@@ -164,12 +185,12 @@ def module(self, module_config):
                 optimizer=self.optimizer,
                 lr_scheduler=self.scheduler,
                 criterion=module.criterion,
-                gradcam=self.config.model.module.gradcam,
+                gradcam=self.gradcam,
             )
         self._module = module
 
     @property
-    def pre_classifier(self) -> torch.nn.Module:
+    def pre_classifier(self) -> nn.Module:
         return self._pre_classifier
 
     @pre_classifier.setter
@@ -178,11 +199,11 @@ def pre_classifier(self, model_config: DictConfig) -> None:
             log.info("Instantiating pre_classifier <%s>", model_config.pre_classifier["_target_"])
             self._pre_classifier = hydra.utils.instantiate(model_config.pre_classifier, _convert_="partial")
         else:
-            log.info("No pre-classifier found in the config: instantiate a torch.nn.Identity instead")
-            self._pre_classifier = torch.nn.Identity()
+            log.info("No pre-classifier found in config: instantiate a torch.nn.Identity instead")
+            self._pre_classifier = nn.Identity()
 
     @property
-    def classifier(self) -> torch.nn.Module:
+    def classifier(self) -> nn.Module:
         return self._classifier
 
     @classifier.setter
@@ -194,7 +215,7 @@ def classifier(self, model_config: DictConfig) -> None:
             raise ValueError("A `classifier` definition must be specified in the config")
 
     @property
-    def model(self) -> torch.nn.Module:
+    def model(self) -> nn.Module:
         return self._model
 
     @model.setter
@@ -227,17 +248,24 @@ def test(self) -> None:
             self.trainer.test(datamodule=self.datamodule, model=self.module, ckpt_path="best")
 
     def export(self) -> None:
-        """Generate deployment model for the task."""
-        if self.export_type is None:
-            raise ValueError("`Export_type` must be specified in the config to export the model")
+        """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
 
-        input_width = self.config.transforms.get("input_width")
-        input_height = self.config.transforms.get("input_height")
+        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"
+            )
+            idx_to_class = {}
+        else:
+            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 runtime model for %s checkpoint", best_model_path)
+        log.info("Saving deployment model for %s checkpoint", best_model_path)
 
         module = self.module.load_from_checkpoint(
             best_model_path,
@@ -245,33 +273,47 @@ def export(self) -> None:
             optimizer=self.optimizer,
             lr_scheduler=self.scheduler,
             criterion=self.module.criterion,
-            gradcam=self.module.gradcam,
+            gradcam=False,
         )
 
-        for export_type in self.export_type:
+        input_shapes = self.export_config.input_shapes
+
+        # TODO: This breaks with bf16 precision!!!
+        half_precision = int(self.trainer.precision) == 16
+
+        for export_type in self.export_config.types:
             if export_type == "torchscript":
-                export_torchscript_model(
-                    module.model,
-                    (1, 3, input_height, input_width),
-                    self.export_folder,
-                    half_precision=int(self.trainer.precision) == 16,
+                out = export_torchscript_model(
+                    model=module.model,
+                    input_shapes=input_shapes,
+                    output_path=self.export_folder,
+                    half_precision=half_precision,
                 )
-        if 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"
-            )
-            idx_to_class = {}
-        else:
-            idx_to_class = {v: k for k, v in self.datamodule.class_to_idx.items()}
-        model_json = {
-            "input_size": [input_width, input_height, 3],
+
+                if out is None:
+                    log.warning("Skipping torchscript export since the model is not supported")
+                    continue
+
+                _, input_shapes = out
+            elif export_type == "pytorch":
+                export_pytorch_model(
+                    model=module.model,
+                    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 = {
+            "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, "model.json"), "w") as f:
-            json.dump(model_json, f)
+        with open(os.path.join(self.export_folder, self.deploy_info_file), "w") as f:
+            json.dump(self.model_json, f)
 
     def generate_report(self) -> None:
         """Generate a report for the task."""
@@ -284,8 +326,12 @@ def generate_report(self) -> None:
             return
 
         log.info("Generating report!")
-
-        predictions_outputs = self.trainer.predict(model=self.module, datamodule=self.datamodule, ckpt_path="best")
+        if not self.run_test or self.config.trainer.get("fast_dev_run"):
+            self.datamodule.setup(stage="test")
+        best_model_path = self.trainer.checkpoint_callback.best_model_path  # type: ignore[union-attr]
+        predictions_outputs = self.trainer.predict(
+            model=self.module, datamodule=self.datamodule, ckpt_path=best_model_path
+        )
         if not predictions_outputs:
             log.warning("There is no prediction to generate the report. Skipping report generation.")
             return
@@ -371,7 +417,8 @@ def generate_report(self) -> None:
                             utils.upload_file_tensorboard(a, tensorboard_logger)
 
     def freeze_layers(self, freeze_parameters_name: List[str]):
-        """Freeze layers specified in freeze_parameters_name
+        """Freeze layers specified in freeze_parameters_name.
+
         Args:
             freeze_parameters_name: Layers that will be frozen during training.
 
@@ -395,7 +442,11 @@ class SklearnClassification(Generic[SklearnClassificationDataModuleT], Task[Skle
         config: The experiment configuration
         device: The device to use. Defaults to None.
         output: Dictionary defining which kind of outputs to generate. Defaults to None.
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
@@ -403,13 +454,13 @@ def __init__(
         config: DictConfig,
         output: DictConfig,
         device: str,
-        export_type: Optional[List[str]] = None,
+        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_type=export_type)
+        super().__init__(config=config, export_config=export_config)
 
         self._device = device
         self.output = output
-        self._backbone: torch.nn.Module
+        self._backbone: nn.Module
         self._trainer: SklearnClassificationTrainer
         self._model: ClassifierMixin
         self.metadata: Dict[str, Any] = {
@@ -418,7 +469,8 @@ def __init__(
             "test_results": [],
             "test_labels": [],
         }
-        self.export_folder: str = "deployment_model"
+        self.export_folder = "deployment_model"
+        self.deploy_info_file = "model.json"
         self.train_dataloader_list: List[torch.utils.data.DataLoader] = []
         self.test_dataloader_list: List[torch.utils.data.DataLoader] = []
 
@@ -433,6 +485,7 @@ def prepare(self) -> None:
         self.backbone = self.config.backbone
 
         self.model = self.config.model
+        # prepare_data() must be explicitly called if the task does not include a lightining training
         self.datamodule.prepare_data()
         self.datamodule.setup(stage="fit")
 
@@ -452,7 +505,7 @@ def model(self, model_config: DictConfig):
         self._model = hydra.utils.instantiate(model_config)
 
     @property
-    def backbone(self) -> torch.nn.Module:
+    def backbone(self) -> nn.Module:
         """Backbone: The backbone."""
         return self._backbone
 
@@ -465,8 +518,10 @@ def backbone(self, backbone_config):
         else:
             log.info("Loading backbone from <%s>", backbone_config.model["_target_"])
             self._backbone = hydra.utils.instantiate(backbone_config.model)
+
+        self._backbone = ModelSignatureWrapper(self._backbone)
         self._backbone.eval()
-        self._backbone = self._backbone.to(self.device)
+        self._backbone.to(self.device)
 
     @property
     def trainer(self) -> SklearnClassificationTrainer:
@@ -497,7 +552,7 @@ def train(self) -> None:
                 raise AttributeError("Cache is only supported when iteration over training is set to 1")
 
             full_dataloader = self.datamodule.full_dataloader()
-            all_features, all_labels = get_feature(
+            all_features, all_labels, _ = get_feature(
                 feature_extractor=self.backbone, dl=full_dataloader, iteration_over_training=1
             )
 
@@ -527,7 +582,7 @@ def train(self) -> None:
                     train_features=all_features_sorted[0:train_len], train_labels=all_labels_sorted[0:train_len]
                 )
 
-                _, pd_cm, accuracy, res = self.trainer.test(
+                _, pd_cm, accuracy, res, _ = self.trainer.test(
                     test_dataloader=test_dataloader,
                     test_features=all_features_sorted[train_len:],
                     test_labels=all_labels_sorted[train_len:],
@@ -537,7 +592,7 @@ def train(self) -> None:
                 )
             else:
                 self.trainer.fit(train_dataloader=train_dataloader)
-                _, pd_cm, accuracy, res = self.trainer.test(
+                _, pd_cm, accuracy, res, _ = self.trainer.test(
                     test_dataloader=test_dataloader,
                     class_to_keep=class_to_keep,
                     idx_to_class=train_dataloader.dataset.idx_to_class,
@@ -581,7 +636,7 @@ def test_full_data(self) -> None:
             log.info("No test data, skipping test")
             return
 
-        _, pd_cm, accuracy, res = self.trainer.test(
+        _, pd_cm, accuracy, res, _ = self.trainer.test(
             test_dataloader=test_dataloader, idx_to_class=idx_to_class, predict_proba=True
         )
 
@@ -601,30 +656,47 @@ def test_full_data(self) -> None:
 
     def export(self) -> None:
         """Generate deployment model for the task."""
-        input_width = self.config.transforms.get("input_width")
-        input_height = self.config.transforms.get("input_height")
-
-        if self.export_type is None:
-            log.info("No export type specified, skipping export")
+        if self.export_config is None or len(self.export_config.types) == 0:
+            log.info("No export type specified skipping export")
             return
 
-        for export_type in self.export_type:
+        input_shapes = self.export_config.input_shapes
+
+        for export_type in self.export_config.types:
             if export_type == "torchscript":
-                export_torchscript_model(
-                    self.backbone, (1, 3, input_height, input_width), self.export_folder, half_precision=False
+                out = export_torchscript_model(
+                    model=self.backbone,
+                    input_shapes=input_shapes,
+                    output_path=self.export_folder,
+                    half_precision=False,
                 )
 
+                if out is None:
+                    log.warning("Skipping torchscript export since the model is not supported")
+                    continue
+
+                _, input_shapes = out
+            elif export_type == "pytorch":
+                os.makedirs(self.export_folder, exist_ok=True)
+                # We only need to save classifier.joblib + backbone config file
+                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)
+
         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_width, input_height, 3],
+            "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, "model.json"), "w") as f:
+        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:
@@ -670,57 +742,56 @@ def execute(self) -> None:
         if self.output.report:
             self.generate_report()
         self.train_full_data()
-        if self.export_type is not None and len(self.export_type) > 0:
+        if self.export_config is not None and len(self.export_config.types) > 0:
             self.export()
         if self.output.test_full_data:
             self.test_full_data()
         self.finalize()
 
 
-class SklearnTestClassification(Task[SklearnClassificationDataModuleT]):
-    """Perform a test of an already trained classification model.
+class SklearnTestClassification(Evaluation[SklearnClassificationDataModuleT]):
+    """Perform a test using an imported SklearnClassification pytorch model.
 
     Args:
         config: The experiment configuration
         output: where to save results
-        experiment_path: path to training experiment generated from SklearnClassification task.
+        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
     """
 
     def __init__(
-        self,
+        self,  # pylint: disable=W0613
         config: DictConfig,
         output: DictConfig,
-        experiment_path: str,
+        model_path: str,
         device: str,
+        gradcam: bool = False,
+        **kwargs: Any,
     ):
-        super().__init__(config=config)
-        self._device = device
+        super().__init__(config=config, model_path=model_path, device=device)
+        self.gradcam = gradcam
         self.output = output
-        self.experiment_path = experiment_path
-        self._backbone: torch.nn.Module
+        self._backbone: nn.Module
         self._classifier: ClassifierMixin
         self.class_to_idx: Dict[str, int]
         self.idx_to_class: Dict[int, str]
-        self.runtime_info_file = "model.json"
         self.test_dataloader: torch.utils.data.DataLoader
         self.metadata: Dict[str, Any] = {
             "test_confusion_matrix": None,
             "test_accuracy": None,
             "test_results": None,
             "test_labels": None,
+            "cams": None,
         }
-        self.model_info: Any
 
     def prepare(self) -> None:
         """Prepare the experiment."""
-        # Read the information of the already trained model
-        with open(os.path.join(self.experiment_path, "deployment_model", self.runtime_info_file), "r") as f:
-            self.model_info = json.load(f)
+        super().prepare()
 
         idx_to_class = {}
         class_to_idx = {}
-        for k, v in self.model_info["classes"].items():
+        for k, v in self.model_data["classes"].items():
             idx_to_class[int(k)] = v
             class_to_idx[v] = int(k)
 
@@ -729,22 +800,28 @@ def prepare(self) -> None:
 
         self.config.datamodule.class_to_idx = class_to_idx
 
-        # Setup datamodule
         self.datamodule = self.config.datamodule
+        # prepare_data() must be explicitly called because there is no lightning training
         self.datamodule.prepare_data()
         self.datamodule.setup(stage="test")
 
         self.test_dataloader = self.datamodule.test_dataloader()
 
-        # TODO: It is not currently possible to load a deployed backbone!!
-        self.backbone = self.config.backbone
-
-        # Load classifier
-        self.classifier = os.path.join(self.experiment_path, "deployment_model", "classifier.joblib")
-
         # Configure trainer
         self.trainer = self.config.trainer
 
+    @property
+    def deployment_model(self):
+        """Deployment model."""
+        return None
+
+    @deployment_model.setter
+    def deployment_model(self, model_path: str):
+        """Set backbone and classifier."""
+        self.backbone = model_path  # type: ignore[assignment]
+        # Load classifier
+        self.classifier = os.path.join(Path(model_path).parent, "classifier.joblib")
+
     @property
     def classifier(self) -> ClassifierMixin:
         """Classifier: The classifier."""
@@ -756,21 +833,32 @@ def classifier(self, classifier_path: str) -> None:
         self._classifier = load(classifier_path)
 
     @property
-    def backbone(self) -> torch.nn.Module:
+    def backbone(self) -> nn.Module:
         """Backbone: The backbone."""
         return self._backbone
 
     @backbone.setter
-    def backbone(self, backbone_config: DictConfig) -> None:
+    def backbone(self, model_path: str) -> None:
         """Load backbone."""
-        if backbone_config.metadata.get("checkpoint"):
-            log.info("Loading backbone from <%s>", backbone_config.metadata.checkpoint)
-            self._backbone = torch.load(backbone_config.metadata.checkpoint)
+        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")
+            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)
+            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("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)
+            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)
 
     @property
     def trainer(self) -> SklearnClassificationTrainer:
@@ -790,8 +878,11 @@ def trainer(self, trainer_config: DictConfig) -> None:
 
     def test(self) -> None:
         """Run the test."""
-        _, pd_cm, accuracy, res = self.trainer.test(
-            test_dataloader=self.test_dataloader, idx_to_class=self.idx_to_class, predict_proba=True
+        _, pd_cm, accuracy, res, cams = self.trainer.test(
+            test_dataloader=self.test_dataloader,
+            idx_to_class=self.idx_to_class,
+            predict_proba=True,
+            gradcam=self.gradcam,
         )
 
         # save results
@@ -801,12 +892,12 @@ def test(self) -> None:
         self.metadata["test_labels"] = [
             self.idx_to_class[i] if i != -1 else "N/A" for i in res["real_label"].unique().tolist()
         ]
+        self.metadata["cams"] = cams
 
     def generate_report(self) -> None:
         """Generate a report for the task."""
         log.info("Generating report!")
         os.makedirs(self.output.folder, exist_ok=True)
-
         save_classification_result(
             results=self.metadata["test_results"],
             output_folder=self.output.folder,
@@ -815,6 +906,7 @@ def generate_report(self) -> None:
             test_dataloader=self.test_dataloader,
             config=self.config,
             output=self.output,
+            grayscale_cams=self.metadata["cams"],
         )
 
     def execute(self) -> None:
@@ -825,18 +917,16 @@ def execute(self) -> None:
             self.generate_report()
         self.finalize()
 
-    @property
-    def device(self) -> str:
-        return self._device
-
 
-class ClassificationEvaluation(Task[ClassificationDataModule]):
-    """Evaluation task for Classification.
+class ClassificationEvaluation(Evaluation[ClassificationDataModuleT]):
+    """Perform a test on an imported Classification pytorch model.
 
     Args:
         config: Task configuration
         output: Configuration for the output
-        model_path: Path to .pt model file
+        model_path: Path to pytorch .pt model file
+        report: Whether to generate the report of the predictions
+        gradcam: Whether to compute gradcams
         device: Device to use for evaluation. If None, the device is automatically determined
 
     """
@@ -846,23 +936,54 @@ def __init__(
         config: DictConfig,
         output: DictConfig,
         model_path: str,
+        report: bool = True,
+        gradcam: bool = False,
         device: Optional[str] = None,
     ):
-        super().__init__(config=config)
-        self.model_data: Dict[str, Any]
-        self._deployment_model: Any
-        self.deployment_model_type: str
-        self.model_path = model_path
-        self.output_path = "test_output"
+        super().__init__(config=config, model_path=model_path, device=device)
+        self.report_path = "test_output"
         self.output = output
-        self.config = config
-        self.report_path = ""
-        self.metadata = {"report_files": []}
-        self.model_info_filename = "model.json"
-        if device is None:
-            self.device = utils.get_device()
+        self.report = report
+        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]
+        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
+
+    @pre_classifier.setter
+    def pre_classifier(self, model_config: DictConfig) -> None:
+        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")
         else:
-            self.device = device
+            log.info("No pre-classifier found in config: instantiate a torch.nn.Identity instead")
+            self._pre_classifier = nn.Identity()
+
+    @property
+    def classifier(self) -> nn.Module:
+        return self._classifier
+
+    @classifier.setter
+    def classifier(self, model_config: DictConfig) -> None:
+        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")
 
     @property
     def deployment_model(self):
@@ -872,55 +993,83 @@ 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)
+        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
+        )
+        if self.gradcam and self.deployment_model_type != "torch":
+            log.warning("To compute gradcams you need to provide the path to an exported .pth state_dict file")
+            self.gradcam = False
 
     def prepare(self) -> None:
         """Prepare the evaluation."""
-        with open(os.path.join(Path(self.model_path).parent, self.model_info_filename)) as f:
-            self.model_data = json.load(f)
-
-        if not isinstance(self.model_data, dict):
-            raise ValueError("Model info file is not a valid json")
-
-        if self.model_data["input_size"][0] != self.config.transforms.input_height:
-            log.warning(
-                f"Input height of the model ({self.model_data['input_size'][0]}) is different from the one specified "
-                + f"in the config ({self.config.transforms.input_height}). Fixing the config."
-            )
-            self.config.transforms.input_height = self.model_data["input_size"][0]
-
-        if self.model_data["input_size"][1] != self.config.transforms.input_width:
-            log.warning(
-                f"Input width of the model ({self.model_data['input_size'][1]}) is different from the one specified "
-                + f"in the config ({self.config.transforms.input_width}). Fixing the config."
-            )
-            self.config.transforms.input_width = self.model_data["input_size"][1]
-
-        self.deployment_model = self.model_path
-
         super().prepare()
+        self.datamodule = self.config.datamodule
         self.datamodule.class_to_idx = {v: int(k) for k, v in self.model_data["classes"].items()}
 
+    def prepare_gradcam(self) -> None:
+        """Initializing gradcam for the predictions."""
+        if isinstance(self.deployment_model.features_extractor, timm.models.resnet.ResNet):
+            target_layers = [
+                cast(BaseNetworkBuilder, self.deployment_model).features_extractor.layer4[-1]  # type: ignore[index]
+            ]
+            self.cam = GradCAM(
+                model=self.deployment_model,
+                target_layers=target_layers,
+                use_cuda=(self.device != "cpu"),
+            )
+            for p in self.deployment_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)
+        else:
+            log.warning("Gradcam not implemented for this backbone, it will not be computed")
+            self.gradcam = False
+
     def test(self) -> None:
         """Perform test."""
         log.info("Running test")
+        # prepare_data() must be explicitly called because there is no training
         self.datamodule.prepare_data()
         self.datamodule.setup(stage="test")
         test_dataloader = self.datamodule.test_dataloader()
 
         image_labels = []
         predicted_classes = []
+        grayscale_cams_list = []
 
-        with torch.no_grad():
+        if self.gradcam:
+            self.prepare_gradcam()
+
+        with torch.set_grad_enabled(self.gradcam):
             for batch_item in tqdm(test_dataloader):
                 im, target = batch_item
-
-                outputs = self.deployment_model(im.to(self.device))
+                im = im.to(self.device).detach()
+                outputs = self.deployment_model(im).detach()
                 probs = torch.softmax(outputs, dim=1)
                 preds = torch.max(probs, dim=1).indices
 
                 predicted_classes.append(preds.tolist())
                 image_labels.extend(target.tolist())
+                if self.gradcam:
+                    with torch.inference_mode(False):
+                        im = im.clone()
+                        if isinstance(self.deployment_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):
+                            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])
+                            zoom_factors = tuple(np.array(new_shape) / np.array(orig_shape))
+                            grayscale_cam = ndimage.zoom(grayscale_cam_low_res, zoom_factors, order=1)
+                            grayscale_cams_list.append(torch.from_numpy(grayscale_cam))
+
+        grayscale_cams: Optional[torch.Tensor] = None
+        if self.gradcam:
+            grayscale_cams = torch.cat(grayscale_cams_list, dim=0)
 
         predicted_classes = [item for sublist in predicted_classes for item in sublist]
         if self.datamodule.class_to_idx is not None:
@@ -938,11 +1087,12 @@ def test(self) -> None:
         self.metadata["test_accuracy"] = test_accuracy
         self.metadata["test_results"] = predicted_classes
         self.metadata["test_labels"] = image_labels
+        self.metadata["grayscale_cams"] = grayscale_cams
 
     def generate_report(self) -> None:
         """Generate a report for the task."""
         log.info("Generating report!")
-        os.makedirs(self.output_path, exist_ok=True)
+        os.makedirs(self.report_path, exist_ok=True)
 
         test_dataset = cast(ImageClassificationListDataset, self.datamodule.test_dataloader().dataset)
         res = pd.DataFrame(
@@ -952,22 +1102,22 @@ def generate_report(self) -> None:
                 "pred_label": self.metadata["test_results"],
             }
         )
-        os.makedirs(self.output_path, exist_ok=True)
+        os.makedirs(self.report_path, exist_ok=True)
         save_classification_result(
             results=res,
-            output_folder=self.output_path,
+            output_folder=self.report_path,
             confmat=self.metadata["test_confusion_matrix"],
             accuracy=self.metadata["test_accuracy"],
             test_dataloader=self.datamodule.test_dataloader(),
             config=self.config,
             output=self.output,
-            grayscale_cams=None,
+            grayscale_cams=self.metadata["grayscale_cams"],
         )
 
     def execute(self) -> None:
         """Execute the evaluation."""
         self.prepare()
         self.test()
-        self.generate_report()
-        log.info("Evaluation finished")
+        if self.report:
+            self.generate_report()
         self.finalize()
diff --git a/quadra/tasks/patch.py b/quadra/tasks/patch.py
index 38a256c0..afb0def6 100644
--- a/quadra/tasks/patch.py
+++ b/quadra/tasks/patch.py
@@ -1,19 +1,21 @@
 import json
 import os
+from pathlib import Path
 from typing import Any, Dict, List, Optional, cast
 
 import hydra
 import torch
 from joblib import dump, load
-from omegaconf import DictConfig
+from omegaconf import DictConfig, OmegaConf
 from sklearn.base import ClassifierMixin
 
 from quadra.datamodules import PatchSklearnClassificationDataModule
 from quadra.datasets.patch import PatchSklearnClassificationTrainDataset
-from quadra.tasks.base import Task
+from quadra.models.base import ModelSignatureWrapper
+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
+from quadra.utils.export import export_torchscript_model, import_deployment_model
 from quadra.utils.patch import RleEncoder, compute_patch_metrics, save_classification_result
 from quadra.utils.patch.dataset import PatchDatasetFileFormat
 
@@ -27,7 +29,11 @@ class PatchSklearnClassification(Task[PatchSklearnClassificationDataModule]):
         config: The experiment configuration
         device: The device to use
         output: Dictionary defining which kind of outputs to generate. Defaults to None.
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
@@ -35,9 +41,9 @@ def __init__(
         config: DictConfig,
         output: DictConfig,
         device: str,
-        export_type: Optional[List[str]] = None,
+        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_type=export_type)
+        super().__init__(config=config, export_config=export_config)
         self.device: str = device
         self.output: DictConfig = output
         self.return_polygon: bool = True
@@ -78,6 +84,8 @@ def backbone(self, backbone_config):
         else:
             log.info("Loading backbone from <%s>", backbone_config.model["_target_"])
             self._backbone = hydra.utils.instantiate(backbone_config.model)
+
+        self._backbone = ModelSignatureWrapper(self._backbone)
         self._backbone.eval()
         self._backbone = self._backbone.to(self.device)
 
@@ -104,6 +112,7 @@ def trainer(self, trainer_config: DictConfig) -> None:
     def train(self) -> None:
         """Train the model."""
         log.info("Starting training...!")
+        # prepare_data() must be explicitly called if the task does not include a lightining training
         self.datamodule.prepare_data()
         self.datamodule.setup(stage="fit")
         class_to_keep = None
@@ -118,7 +127,7 @@ def train(self) -> None:
         val_dataloader = self.datamodule.val_dataloader()
         train_dataset = cast(PatchSklearnClassificationTrainDataset, train_dataloader.dataset)
         self.trainer.fit(train_dataloader=train_dataloader)
-        _, pd_cm, accuracy, res = self.trainer.test(
+        _, pd_cm, accuracy, res, _ = self.trainer.test(
             test_dataloader=val_dataloader,
             class_to_keep=class_to_keep,
             idx_to_class=train_dataset.idx_to_class,
@@ -198,18 +207,29 @@ def generate_report(self) -> None:
 
     def export(self) -> None:
         """Generate deployment model for the task."""
-        if self.export_type is None or len(self.export_type) == 0:
+        if self.export_config is None or len(self.export_config.types) == 0:
             log.info("No export type specified skipping export")
             return
-        input_width = self.config.transforms.get("input_width")
-        input_height = self.config.transforms.get("input_height")
 
-        for export_type in self.export_type:
+        os.makedirs(self.export_folder, exist_ok=True)
+
+        input_shapes = self.export_config.input_shapes
+
+        for export_type in self.export_config.types:
             if export_type == "torchscript":
-                export_torchscript_model(
-                    self.backbone, (1, 3, input_height, input_width), self.export_folder, half_precision=False
+                out = export_torchscript_model(
+                    model=self.backbone,
+                    input_shapes=input_shapes,
+                    output_path=self.export_folder,
+                    half_precision=False,
                 )
 
+                if out is None:
+                    log.warning("Skipping torchscript export since the model is not supported")
+                    continue
+
+                _, input_shapes = out
+
         dump(self.model, os.path.join(self.export_folder, "classifier.joblib"))
 
         dataset_info = self.datamodule.info
@@ -223,7 +243,7 @@ def export(self) -> None:
         idx_to_class = {v: k for k, v in self.datamodule.class_to_idx.items()}
 
         model_json = {
-            "input_size": [input_width, input_height, 3],
+            "input_size": input_shapes,
             "classes": idx_to_class,
             "mean": self.config.transforms.mean,
             "std": self.config.transforms.std,
@@ -245,18 +265,18 @@ def execute(self) -> None:
         self.train()
         if self.output.report:
             self.generate_report()
-        if self.export_type is not None and len(self.export_type) > 0:
+        if self.export_config is not None and len(self.export_config.types) > 0:
             self.export()
         self.finalize()
 
 
-class PatchSklearnTestClassification(Task[PatchSklearnClassificationDataModule]):
+class PatchSklearnTestClassification(Evaluation[PatchSklearnClassificationDataModule]):
     """Perform a test of an already trained classification model.
 
     Args:
         config: The experiment configuration
         output: where to save resultss
-        experiment_path: path to training experiment generated from SklearnClassification task.
+        model_path: path to trained model from PatchSklearnClassification task.
         device: the device where to run the model (cuda or cpu). Defaults to 'cpu'.
     """
 
@@ -264,69 +284,56 @@ def __init__(
         self,
         config: DictConfig,
         output: DictConfig,
-        experiment_path: str,
+        model_path: str,
         device: str = "cpu",
     ):
-        super().__init__(config=config)
-        self._device = device
+        super().__init__(config=config, model_path=model_path, device=device)
         self.output = output
-        self.experiment_path = experiment_path
         self._backbone: torch.nn.Module
         self._classifier: ClassifierMixin
         self.class_to_idx: Dict[str, int]
         self.idx_to_class: Dict[int, str]
-        self.runtime_info_file = "model.json"
         self.metadata: Dict[str, Any] = {
             "test_confusion_matrix": None,
             "test_accuracy": None,
             "test_results": None,
             "test_labels": None,
         }
-        self.model_info: Any
         self.class_to_skip: List[str] = []
         self.reconstruction_results: Dict[str, Any]
         self.return_polygon: bool = True
 
     def prepare(self) -> None:
         """Prepare the experiment."""
-        # Read the information of the already trained model
-        with open(os.path.join(self.experiment_path, "deployment_model", self.runtime_info_file), "r") as f:
-            self.model_info = json.load(f)
+        super().prepare()
 
         idx_to_class = {}
         class_to_idx = {}
-        for k, v in self.model_info["classes"].items():
+        for k, v in self.model_data["classes"].items():
             idx_to_class[int(k)] = v
             class_to_idx[v] = int(k)
 
         self.idx_to_class = idx_to_class
         self.class_to_idx = class_to_idx
-
         self.config.datamodule.class_to_idx = class_to_idx
 
-        # Setup datamodule
         self.datamodule = self.config.datamodule
-
-        self.backbone = self.config.backbone
-
-        # Load classifier
-        self.classifier = os.path.join(self.experiment_path, "deployment_model", "classifier.joblib")
-
         # Configure trainer
         self.trainer = self.config.trainer
 
     def test(self) -> None:
         """Run the test."""
+        # prepare_data() must be explicitly called because there is no lightning training
         self.datamodule.prepare_data()
         self.datamodule.setup(stage="test")
         test_dataloader = self.datamodule.test_dataloader()
 
-        self.class_to_skip = self.model_info["class_to_skip"] if hasattr(self.model_info, "class_to_skip") else None
+        self.class_to_skip = self.model_data["class_to_skip"] if hasattr(self.model_data, "class_to_skip") else None
         class_to_keep = None
 
         if self.class_to_skip is not None:
             class_to_keep = [x for x in self.datamodule.class_to_idx.keys() if x not in self.class_to_skip]
-        _, pd_cm, accuracy, res = self.trainer.test(
+        _, pd_cm, accuracy, res, _ = self.trainer.test(
             test_dataloader=test_dataloader,
             idx_to_class=self.idx_to_class,
             predict_proba=True,
@@ -341,6 +348,18 @@ def test(self) -> None:
             self.idx_to_class[i] if i != -1 else "N/A" for i in res["real_label"].unique().tolist()
         ]
 
+    @property
+    def deployment_model(self):
+        """Deployment model."""
+        return None
+
+    @deployment_model.setter
+    def deployment_model(self, model_path: str):
+        """Set backbone and classifier."""
+        self.backbone = model_path  # type: ignore[assignment]
+        # Load classifier
+        self.classifier = os.path.join(Path(model_path).parent, "classifier.joblib")
+
     @property
     def classifier(self) -> ClassifierMixin:
         """Classifier: The classifier."""
@@ -357,16 +376,27 @@ def backbone(self) -> torch.nn.Module:
         return self._backbone
 
     @backbone.setter
-    def backbone(self, backbone_config: DictConfig) -> None:
+    def backbone(self, model_path: str) -> None:
         """Load backbone."""
-        if backbone_config.metadata.get("checkpoint"):
-            log.info("Loading backbone from <%s>", backbone_config.metadata.checkpoint)
-            self._backbone = torch.load(backbone_config.metadata.checkpoint)
+        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")
+            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)
+            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("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)
+            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)
 
     @property
     def trainer(self) -> SklearnClassificationTrainer:
@@ -454,7 +484,3 @@ def execute(self) -> None:
         if self.output.report:
             self.generate_report()
         self.finalize()
-
-    @property
-    def device(self) -> str:
-        return self._device
diff --git a/quadra/tasks/segmentation.py b/quadra/tasks/segmentation.py
index 4ee41afd..11de316c 100644
--- a/quadra/tasks/segmentation.py
+++ b/quadra/tasks/segmentation.py
@@ -11,6 +11,7 @@
 
 from quadra.callbacks.mlflow import get_mlflow_logger
 from quadra.datamodules import SegmentationDataModule, SegmentationMulticlassDataModule
+from quadra.models.base import ModelSignatureWrapper
 from quadra.modules.base import SegmentationModel
 from quadra.tasks.base import Evaluation, LightningTask
 from quadra.utils import utils
@@ -34,7 +35,11 @@ class Segmentation(Generic[SegmentationDataModuleT], LightningTask[SegmentationD
         run_test: If True, run test after training. Defaults to False.
         evaluate: Dict with evaluation parameters. Defaults to None.
         report: If True, create report after training. Defaults to False.
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
@@ -45,29 +50,34 @@ def __init__(
         run_test: bool = False,
         evaluate: Optional[DictConfig] = None,
         report: bool = False,
-        export_type: Optional[List[str]] = None,
+        export_config: Optional[DictConfig] = None,
     ):
         super().__init__(
             config=config,
             checkpoint_path=checkpoint_path,
             run_test=run_test,
             report=report,
-            export_type=export_type,
+            export_config=export_config,
         )
         self.evaluate = evaluate
         self.num_viz_samples = num_viz_samples
         self.export_folder: str = "deployment_model"
         self.exported_model_path: Optional[str] = None
         if self.evaluate and any(self.evaluate.values()):
-            if self.export_type is None or len(self.export_type) == 0 or "torchscript" not in self.export_type:
+            if (
+                self.export_config is None
+                or len(self.export_config.types) == 0
+                or "torchscript" not in self.export_config.types
+            ):
                 log.info(
                     "Evaluation is enabled, but training does not export a deployment model. Automatically export the "
                     "model as torchscript."
                 )
-                if self.export_type is None:
-                    self.export_type = ["torchscript"]
+                if self.export_config is None:
+                    self.export_config = DictConfig({"types": ["torchscript"]})
                 else:
-                    self.export_type.append("torchscript")
+                    self.export_config.types.append("torchscript")
+
             if not self.report:
                 log.info("Evaluation is enabled, but reporting is disabled. Enabling reporting automatically.")
                 self.report = True
@@ -91,6 +101,7 @@ def module(self, module_config) -> None:
                 module_config.model.num_classes = len(self.datamodule.idx_to_class) + 1
 
         model = hydra.utils.instantiate(module_config.model)
+        model = ModelSignatureWrapper(model)
         log.info("Instantiating optimizer <%s>", self.config.optimizer["_target_"])
         param_list = []
         for param in model.parameters():
@@ -113,9 +124,6 @@ def prepare(self) -> None:
     def export(self) -> None:
         """Generate a deployment model for the task."""
         log.info("Exporting model ready for deployment")
-        input_width = self.config.transforms.get("input_width")
-        input_height = self.config.transforms.get("input_height")
-
         # Get best model!
         if self.trainer.checkpoint_callback is None:
             raise ValueError("No checkpoint callback found in the trainer")
@@ -136,19 +144,37 @@ def export(self) -> None:
         else:
             log.info("idx_to_class is present")
             classes = self.config.datamodule.idx_to_class
-        if self.export_type is None:
+
+        if self.export_config is None:
             raise ValueError(
                 "No export type specified. This should not happen, please check if you have set "
                 "the export_type or assign it to a default value."
             )
-        for export_type in self.export_type:
+
+        input_shapes = self.export_config.input_shapes
+
+        half_precision = self.trainer.precision == 16
+
+        for export_type in self.export_config.types:
             if export_type == "torchscript":
-                self.exported_model_path = export_torchscript_model(
-                    module.model, (1, 3, input_height, input_width), self.export_folder, half_precision=False
+                out = export_torchscript_model(
+                    model=module.model,
+                    input_shapes=input_shapes,
+                    output_path=self.export_folder,
+                    half_precision=half_precision,
                 )
 
+                if out is None:
+                    log.warning("Skipping torchscript export since the model is not supported")
+                    continue
+
+                self.exported_model_path, input_shapes = out
+
+        if input_shapes is None:
+            log.warning("Not able to export the model in any format")
+
         model_json = {
-            "input_size": [input_width, input_height, 3],
+            "input_size": input_shapes,
             "classes": classes,
             "mean": self.config.transforms.mean,
             "std": self.config.transforms.std,
@@ -183,7 +209,7 @@ def generate_report(self) -> None:
                     for task in eval_tasks:
                         for file in task.metadata["report_files"]:
                             mflow_logger.experiment.log_artifact(
-                                run_id=mflow_logger.run_id, local_path=file, artifact_path=task.report_folder
+                                run_id=mflow_logger.run_id, local_path=file, artifact_path=task.report_path
                             )
 
                 if tensorboard_logger is not None and self.config.core.get("upload_artifacts"):
@@ -206,13 +232,13 @@ def generate_report(self) -> None:
                                 utils.upload_file_tensorboard(file, tensorboard_logger)
 
 
-class SegmentationEvaluation(Evaluation):
+class SegmentationEvaluation(Evaluation[SegmentationDataModuleT]):
     """Segmentation Evaluation Task with deployment models.
 
     Args:
         config: The experiment configuration
         model_path: The experiment path.
-        report_folder: The report folder. Defaults to None.
+        device: Device to use for evaluation. If None, the device is automatically determined.
 
     Raises:
         ValueError: If the model path is not provided
@@ -222,18 +248,28 @@ def __init__(
         self,
         config: DictConfig,
         model_path: str,
-        report_folder: Optional[str] = None,
+        device: Optional[str] = "cpu",
     ):
-        super().__init__(config=config, model_path=model_path, report_folder=report_folder)
+        super().__init__(config=config, model_path=model_path, device=device)
         self.config = config
-        self.metadata = {"report_files": []}
-
-        # TODO: It's not possible to specify the device from outside!!
-        self.device = utils.get_device(config.trainer.accelerator != "cpu")
 
     def save_config(self) -> None:
         """Skip saving the config."""
 
+    def prepare(self) -> None:
+        """Prepare the evaluation."""
+        super().prepare()
+        # TODO: Why we propagate mean and std only in Segmentation?
+        self.config.transforms.mean = self.model_data["mean"]
+        self.config.transforms.std = self.model_data["std"]
+        # Setup datamodule
+        if hasattr(self.config.datamodule, "idx_to_class"):
+            idx_to_class = self.model_data["classes"]  # dict {index: class}
+            self.config.datamodule.idx_to_class = idx_to_class
+        self.datamodule = self.config.datamodule
+        # prepare_data() must be explicitly called because there is no lightning training
+        self.datamodule.prepare_data()
+
     @torch.no_grad()
     def inference(
         self, dataloader: DataLoader, deployment_model: torch.nn.Module, device: torch.device
@@ -269,18 +305,18 @@ class SegmentationAnalysisEvaluation(SegmentationEvaluation):
     Args:
         config: The experiment configuration
         model_path: The model path.
-        report_folder: The report folder. Defaults to "analysis_report".
+        device: Device to use for evaluation. If None, the device is automatically determined.
     """
 
-    def __init__(self, config: DictConfig, model_path: str, report_folder: str = "analysis_report"):
-        super().__init__(config=config, model_path=model_path, report_folder=report_folder)
+    def __init__(
+        self,
+        config: DictConfig,
+        model_path: str,
+        device: Optional[str] = "cpu",
+    ):
+        super().__init__(config=config, model_path=model_path, device=device)
         self.test_output: Dict[str, Any] = {}
 
-    def prepare(self) -> None:
-        """Prepare the experiment."""
-        super().prepare()
-        self.datamodule = self.config.datamodule
-
     def train(self) -> None:
         """Skip training."""
 
diff --git a/quadra/tasks/ssl.py b/quadra/tasks/ssl.py
index 290e8dd2..84cd3af8 100644
--- a/quadra/tasks/ssl.py
+++ b/quadra/tasks/ssl.py
@@ -1,6 +1,6 @@
 import json
 import os
-from typing import Any, List, Optional, Tuple, cast
+from typing import Any, List, Optional, Tuple, Union, cast
 
 import hydra
 import torch
@@ -13,6 +13,7 @@
 from tqdm import tqdm
 
 from quadra.callbacks.scheduler import WarmupInit
+from quadra.models.base import ModelSignatureWrapper
 from quadra.tasks.base import LightningTask, Task
 from quadra.utils import utils
 from quadra.utils.export import export_torchscript_model, import_deployment_model
@@ -28,7 +29,11 @@ class SSL(LightningTask):
         checkpoint_path: The path to the checkpoint to load the model from Defaults to None
         report: Whether to create the report
         run_test: Whether to run final test
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
@@ -37,14 +42,14 @@ def __init__(
         run_test: bool = False,
         report: bool = False,
         checkpoint_path: Optional[str] = None,
-        export_type: Optional[List[str]] = None,
+        export_config: Optional[DictConfig] = None,
     ):
         super().__init__(
             config=config,
             checkpoint_path=checkpoint_path,
             run_test=run_test,
             report=report,
-            export_type=export_type,
+            export_config=export_config,
         )
         self._backbone: nn.Module
         self._optimizer: torch.optim.Optimizer
@@ -95,29 +100,33 @@ def test(self) -> None:
 
     def export(self) -> None:
         """Deploy a model ready for production."""
-        if self.export_type is None or len(self.export_type) == 0:
+        if self.export_config is None or len(self.export_config.types) == 0:
             log.info("No export type specified skipping export")
             return
 
-        if self.export_type is None:
-            return
-
-        input_width = self.config.transforms.get("input_width")
-        input_height = self.config.transforms.get("input_height")
         mean = self.config.transforms.mean
         std = self.config.transforms.std
-        half_precision = self.trainer.precision == 16
-
-        if "torchscript" in self.export_type:
-            export_torchscript_model(
-                model=cast(nn.Module, self.module.model),
-                input_shape=(1, 3, input_height, input_width),
-                half_precision=half_precision,
-                output_path=self.export_folder,
-            )
+        half_precision = int(self.trainer.precision) == 16
+
+        input_shapes = self.export_config.input_shapes
+
+        for export_type in self.export_config.types:
+            if export_type == "torchscript":
+                out = export_torchscript_model(
+                    model=cast(nn.Module, self.module.model),
+                    input_shapes=input_shapes,
+                    output_path=self.export_folder,
+                    half_precision=half_precision,
+                )
+
+                if out is None:
+                    log.warning("Skipping torchscript export since the model is not supported")
+                    continue
+
+                _, input_shapes = out
 
         model_json = {
-            "input_size": [3, input_width, input_height],
+            "input_size": input_shapes,
             "classes": None,
             "mean": mean,
             "std": std,
@@ -132,21 +141,25 @@ class Simsiam(SSL):
 
     Args:
         config: the main config
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
         checkpoint_path: if a checkpoint is specified, then it will return a trained model,
             with weights loaded from the checkpoint path specified.
             Defaults to None.
         run_test: Whether to run final test
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
         self,
         config: DictConfig,
-        export_type: Optional[List[str]] = None,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
+        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_type=export_type, checkpoint_path=checkpoint_path, run_test=run_test)
+        super().__init__(config=config, export_config=export_config, checkpoint_path=checkpoint_path, run_test=run_test)
         self.backbone: nn.Module
         self.projection_mlp: nn.Module
         self.prediction_mlp: nn.Module
@@ -203,21 +216,25 @@ class SimCLR(SSL):
 
     Args:
         config: the main config
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
         checkpoint_path: if a checkpoint is specified, then it will return a trained model,
             with weights loaded from the checkpoint path specified.
             Defaults to None.
         run_test: Whether to run final test
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
         self,
         config: DictConfig,
-        export_type: Optional[List[str]] = None,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
+        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_type=export_type, checkpoint_path=checkpoint_path, run_test=run_test)
+        super().__init__(config=config, export_config=export_config, checkpoint_path=checkpoint_path, run_test=run_test)
         self.backbone: nn.Module
         self.projection_mlp: nn.Module
 
@@ -258,6 +275,7 @@ def module(self, module_config):
                 lr_scheduler=self.scheduler,
             )
         self._module = module
+        self._module.model = ModelSignatureWrapper(self._module.model)
 
 
 class Barlow(SimCLR):
@@ -265,22 +283,25 @@ class Barlow(SimCLR):
 
     Args:
         config: the main config
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
         checkpoint_path: if a checkpoint is specified, then it will return a trained model,
             with weights loaded from the checkpoint path specified.
             Defaults to None.
         run_test: Whether to run final test
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
 
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
         self,
         config: DictConfig,
-        export_type: Optional[List[str]] = None,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
+        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_type=export_type, checkpoint_path=checkpoint_path, run_test=run_test)
+        super().__init__(config=config, export_config=export_config, checkpoint_path=checkpoint_path, run_test=run_test)
 
     def prepare(self) -> None:
         """Prepare the experiment."""
@@ -305,26 +326,29 @@ class BYOL(SSL):
 
     Args:
         config: the main config
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
         checkpoint_path: if a checkpoint is specified, then it will return a trained model,
             with weights loaded from the checkpoint path specified.
             Defaults to None.
         run_test: Whether to run final test
-        **kwargs: Keyword arguments
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
 
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
+        **kwargs: Keyword arguments
     """
 
     def __init__(
         self,
         config: DictConfig,
-        export_type: Optional[List[str]] = None,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
+        export_config: Optional[DictConfig] = None,
         **kwargs: Any,
     ):
         super().__init__(
             config=config,
-            export_type=export_type,
+            export_config=export_config,
             checkpoint_path=checkpoint_path,
             run_test=run_test,
             **kwargs,
@@ -392,21 +416,25 @@ class DINO(SSL):
 
     Args:
         config: the main config
-        export_type: List of export method for the model, e.g. [torchscript]. Defaults to None.
         checkpoint_path: if a checkpoint is specified, then it will return a trained model,
             with weights loaded from the checkpoint path specified.
             Defaults to None.
         run_test: Whether to run final test
+        export_config: Dictionary containing the export configuration, it should contain the following keys:
+
+            - `types`: List of types to export.
+            - `input_shapes`: Optional list of input shapes to use, they must be in the same order of the forward
+                arguments.
     """
 
     def __init__(
         self,
         config: DictConfig,
-        export_type: Optional[List[str]] = None,
         checkpoint_path: Optional[str] = None,
         run_test: bool = False,
+        export_config: Optional[DictConfig] = None,
     ):
-        super().__init__(config=config, export_type=export_type, checkpoint_path=checkpoint_path, run_test=run_test)
+        super().__init__(config=config, export_config=export_config, checkpoint_path=checkpoint_path, run_test=run_test)
         self.student_model: nn.Module
         self.teacher_model: nn.Module
         self.student_projection_mlp: nn.Module
@@ -501,7 +529,7 @@ def __init__(
         self.embedding_writer = SummaryWriter(self.embeddings_path)
         self.writer_step = 0  # for tensorboard
         self.embedding_image_size = embedding_image_size
-        self._deployment_model: RecursiveScriptModule
+        self._deployment_model: Union[RecursiveScriptModule, nn.Module]
         self.deployment_model_type: str
 
     @property
diff --git a/quadra/trainers/classification.py b/quadra/trainers/classification.py
index 5c7109f0..cca99515 100644
--- a/quadra/trainers/classification.py
+++ b/quadra/trainers/classification.py
@@ -69,10 +69,11 @@ def fit(
 
         if train_dataloader is not None:  # train_features is None or train_labels is None:
             log.info("Extracting features from training set")
-            train_features, train_labels = get_feature(
+            train_features, train_labels, _ = get_feature(
                 feature_extractor=self.backbone,
                 dl=train_dataloader,
                 iteration_over_training=self.iteration_over_training,
+                gradcam=False,
             )
         else:
             log.info("Using cached features for training set")
@@ -95,7 +96,11 @@ def test(
         class_to_keep: Optional[List[int]] = None,
         idx_to_class: Optional[Dict[int, str]] = None,
         predict_proba: bool = True,
-    ) -> Union[Tuple[Union[str, Dict], DataFrame, float, DataFrame], Tuple[None, None, None, DataFrame]]:
+        gradcam: bool = False,
+    ) -> Union[
+        Tuple[Union[str, Dict], DataFrame, float, DataFrame, Optional[np.ndarray]],
+        Tuple[None, None, None, DataFrame, Optional[np.ndarray]],
+    ]:
         """Test classifier on test set.
 
         Args:
@@ -105,11 +110,26 @@ def test(
             class_to_keep: list of class to keep
             idx_to_class: dictionary mapping class index to class name
             predict_proba: if True, predict also probability for each test image
+            gradcam: Whether to compute gradcam
+
+        Returns:
+            cl_rep: Classification report
+            pd_cm: Confusion matrix dataframe
+            accuracy: Test accuracy
+            res: Test results
+            cams: Gradcams
         """
+        cams = None
         # Extract feature
         if test_features is None:
             log.info("Extracting features from test set")
-            test_features, final_test_labels = get_feature(feature_extractor=self.backbone, dl=test_dataloader)
+            test_features, final_test_labels, cams = get_feature(
+                feature_extractor=self.backbone,
+                dl=test_dataloader,
+                gradcam=gradcam,
+                classifier=self.classifier,
+                input_shape=(self.input_shape[2], self.input_shape[0], self.input_shape[1]),
+            )
         else:
             if test_labels is None:
                 raise ValueError("Test labels must be provided when using cached data")
@@ -143,6 +163,8 @@ def test(
 
         if not all(t == -1 for t in filtered_test_labels):
             test_real_label_cm = np.array(filtered_test_labels)
+            if cams is not None:
+                cams = cams[test_real_label_cm != -1]  # TODO: Is class_to_keep still used?
             pred_labels_cm = np.array(test_prediction_label)[test_real_label_cm != -1]
             test_real_label_cm = test_real_label_cm[test_real_label_cm != -1].astype(pred_labels_cm.dtype)
             cl_rep, pd_cm, accuracy = get_results(test_real_label_cm, pred_labels_cm, idx_to_class)
@@ -150,6 +172,6 @@ def test(
             if predict_proba:
                 res["probability"] = test_probability
 
-            return cl_rep, pd_cm, accuracy, res
+            return cl_rep, pd_cm, accuracy, res, cams
 
-        return None, None, None, res
+        return None, None, None, res, cams
diff --git a/quadra/utils/classification.py b/quadra/utils/classification.py
index 445c9723..27064681 100644
--- a/quadra/utils/classification.py
+++ b/quadra/utils/classification.py
@@ -422,6 +422,8 @@ def save_classification_result(
     # Save csv
     results.to_csv(os.path.join(output_folder, "test_results.csv"), index_label="index")
     if grayscale_cams is None:
+        log.info("Plotting only original examples, set gradcam = true in config file to also plot gradcam examples")
+
         save_gradcams = False
     else:
         log.info("Plotting original and gradcam examples")
@@ -501,7 +503,7 @@ def save_classification_result(
             )
 
     else:
-        log.info("Not generating discordant/concordant examples. Check task:output:example in config file")
+        log.info("Not generating discordant/concordant examples. Check task.output.example in config file")
 
 
 def get_results(
diff --git a/quadra/utils/export.py b/quadra/utils/export.py
index b5a54bdf..c66040c4 100644
--- a/quadra/utils/export.py
+++ b/quadra/utils/export.py
@@ -1,50 +1,93 @@
 import os
-from typing import Optional, Tuple, cast
+from typing import Any, List, Optional, Sequence, Tuple, Union, cast
 
 import torch
 from anomalib.models.cflow import CflowLightning
 from torch import nn
 from torch.jit._script import RecursiveScriptModule
 
-from quadra.utils.utils import get_logger
+from quadra.models.base import ModelSignatureWrapper
 
-log = get_logger(__name__)
+# TODO: Solve circular import as it is not possible to import get_logger right now
+
+
+def generate_torch_inputs(
+    input_shapes: List[Any], device: str, half_precision: bool = False, dtype: torch.dtype = torch.float32
+) -> Union[List[Any], Tuple[Any, ...], torch.Tensor]:
+    """Given a list of input shapes that can contain either lists, tuples or dicts, with tuples being the input shapes
+    of the model, generate a list of torch tensors with the given device and dtype.
+    """
+    if isinstance(input_shapes, list):
+        if any(isinstance(inp, (Sequence, dict)) for inp in input_shapes):
+            return [generate_torch_inputs(inp, device, half_precision, dtype) for inp in input_shapes]
+
+        # Base case
+        inp = torch.randn((1, *input_shapes), dtype=dtype, device=device)
+
+    if isinstance(input_shapes, dict):
+        return {k: generate_torch_inputs(v, device, half_precision, dtype) for k, v in input_shapes.items()}
+
+    if isinstance(input_shapes, tuple):
+        if any(isinstance(inp, (Sequence, dict)) for inp in input_shapes):
+            # The tuple contains a list, tuple or dict
+            return tuple(generate_torch_inputs(inp, device, half_precision, dtype) for inp in input_shapes)
+
+        # Base case
+        inp = torch.randn((1, *input_shapes), dtype=dtype, device=device)
+
+    if half_precision:
+        inp = inp.half()
+
+    return inp
 
 
 def export_torchscript_model(
     model: nn.Module,
-    input_shape: Tuple[int, ...],
     output_path: str,
+    input_shapes: Optional[List[Any]] = None,
     half_precision: bool = False,
     model_name: str = "model.pt",
-) -> Optional[str]:
+) -> Optional[Tuple[str, Any]]:
     """Export a PyTorch model with TorchScript.
 
     Args:
         model: PyTorch model to be exported
-        input_shape: Input shape for tracing
+        input_shapes: Inputs shape for tracing
         output_path: Path to save the model
         half_precision: If True, the model will be exported with half precision
         model_name: Name of the exported model
 
     Returns:
-        If the model is exported successfully, the path to the model is returned.
+        If the model is exported successfully, the path to the model and the input shape are returned.
 
     """
-    model.eval()
+    if isinstance(model, ModelSignatureWrapper):
+        if input_shapes is None:
+            input_shapes = model.input_shapes
+        model = model.instance
+
+    if input_shapes is None:
+        # log.warning(
+        #    "Input shape is None, can not trace model! Please provide input_shapes in the task export configuration."
+        # )
+        return None
+
     if isinstance(model, CflowLightning):
-        log.warning("Exporting cflow model with torchscript is not supported yet.")
+        # log.warning("Exporting cflow model with torchscript is not supported yet.")
         return None
 
+    model.eval()
     if half_precision:
-        log.info("Jitting model with half precision on GPU")
+        # log.info("Jitting model with half precision on GPU")
         model.to("cuda:0")
         model = model.half()
-        inp = torch.randn(input_shape, dtype=torch.float16, device="cuda:0")
+        inp = generate_torch_inputs(
+            input_shapes=input_shapes, device="cuda:0", half_precision=True, dtype=torch.float16
+        )
     else:
-        log.info("Jitting model with double precision")
+        # log.info("Jitting model with double precision")
         model.cpu()
-        inp = torch.randn(input_shape, dtype=torch.float32)
+        inp = generate_torch_inputs(input_shapes=input_shapes, device="cpu", half_precision=False, dtype=torch.float32)
 
     with torch.no_grad():
         model_jit = torch.jit.trace(model, inp)
@@ -54,27 +97,64 @@ def export_torchscript_model(
     model_path = os.path.join(output_path, model_name)
     model_jit.save(model_path)
 
-    log.info("Torchscript model saved to %s", os.path.join(os.getcwd(), model_path))
+    # log.info("Torchscript model saved to %s", os.path.join(os.getcwd(), model_path))
+
+    return os.path.join(os.getcwd(), model_path), input_shapes
+
+
+def export_pytorch_model(model: nn.Module, output_path: str, model_name: str = "model.pth") -> str:
+    """Export pytorch model's parameter dictionary using a deserialized state_dict.
+
+    Args:
+        model: PyTorch model to be exported
+        output_path: Path to save the model
+        model_name: Name of the exported model
+
+    Returns:
+        If the model is exported successfully, the path to the model is returned.
+
+    """
+    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))
 
     return os.path.join(os.getcwd(), model_path)
 
 
 # TODO: Update signature when new models are added
-def import_deployment_model(model_path: str, device: str) -> Tuple[RecursiveScriptModule, str]:
-    """Try to import a model for deployment, currently only supports torchscript .pt files.
+def import_deployment_model(
+    model_path: str, device: str, model: Optional[nn.Module] = None
+) -> Tuple[Union[RecursiveScriptModule, nn.Module], str]:
+    """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
         device: Device to load the model on
+        model: Pytorch model needed to load the parameter dictionary
 
     Returns:
         A tuple containing the model and the model type
     """
+
     file_extension = os.path.splitext(os.path.basename(model_path))[1]
     if file_extension == ".pt":
         model = cast(RecursiveScriptModule, torch.jit.load(model_path))
         model.eval()
         model.to(device)
+
         return model, "torchscript"
+    if file_extension == ".pth":
+        if model is None:
+            raise ValueError("Model is not defined, can not load state_dict!")
+
+        model.load_state_dict(torch.load(model_path))
+        model.eval()
+        model.to(device)
+
+        return model, "torch"
 
-    raise ValueError(f"Unable to load model with extension {file_extension}, valid extensions are: ['.pt']")
+    raise ValueError(f"Unable to load model with extension {file_extension}, valid extensions are: ['.pt', 'pth']")
diff --git a/quadra/utils/mlflow.py b/quadra/utils/mlflow.py
new file mode 100644
index 00000000..438e1f1e
--- /dev/null
+++ b/quadra/utils/mlflow.py
@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+try:
+    from mlflow.models import infer_signature  # noqa
+    from mlflow.models.signature import ModelSignature  # noqa
+
+    MLFLOW_AVAILABLE = True
+except ImportError:
+    MLFLOW_AVAILABLE = False
+
+from typing import Any, Sequence, TypeVar
+
+import torch
+from pytorch_lightning import Trainer
+from pytorch_lightning.loggers import MLFlowLogger
+from torch import nn
+
+NnModuleT = TypeVar("NnModuleT", bound=nn.Module)
+
+
+@torch.inference_mode()
+def infer_signature_torch_model(model: NnModuleT, data: list[Any]) -> ModelSignature | None:
+    """Infer input and output signature for a PyTorch/Torchscript model."""
+    model = model.eval()
+    model = model.cpu()
+    model_output = model(*data)
+
+    try:
+        output_signature = infer_signature_input_torch(model_output)
+
+        if len(data) == 1:
+            signature_input = infer_signature_input_torch(data[0])
+        else:
+            signature_input = infer_signature_input_torch(data)
+    except ValueError:
+        # TODO: Solve circular import as it is not possible to import get_logger right now
+        # log.warning("Unable to infer signature for model output type %s", type(model_output))
+        return None
+
+    return infer_signature(signature_input, output_signature)
+
+
+def infer_signature_input_torch(input_tensor: Any) -> Any:
+    """Recursively infer the signature input format to pass to mlflow.models.infer_signature.
+
+    Raises:
+        ValueError: If the input type is not supported or when nested dicts or sequences are encountered.
+    """
+    if isinstance(input_tensor, Sequence):
+        # Mlflow currently does not support sequence outputs, so we use a dict instead
+        signature = {}
+        for i, x in enumerate(input_tensor):
+            if isinstance(x, Sequence):
+                # Nested signature is currently not supported by mlflow
+                raise ValueError("Nested sequences are not supported")
+                # TODO: Enable this once mlflow supports nested signatures
+                # signature[f"output_{i}"] = {f"output_{j}": infer_signature_torch(y) for j, y in enumerate(x)}
+            if isinstance(x, dict):
+                # Nested dicts are not supported
+                raise ValueError("Nested dicts are not supported")
+
+            signature[f"output_{i}"] = infer_signature_input_torch(x)
+    elif isinstance(input_tensor, torch.Tensor):
+        signature = input_tensor.cpu().numpy()
+    elif isinstance(input_tensor, dict):
+        signature = {}
+        for k, v in input_tensor.items():
+            if isinstance(v, dict):
+                # Nested dicts are not supported
+                raise ValueError("Nested dicts are not supported")
+            if isinstance(v, Sequence):
+                # Nested signature is currently not supported by mlflow
+                raise ValueError("Nested sequences are not supported")
+
+            signature[k] = infer_signature_input_torch(v)
+    else:
+        raise ValueError(f"Unable to infer signature for model output type {type(input_tensor)}")
+
+    return signature
+
+
+def get_mlflow_logger(trainer: Trainer) -> MLFlowLogger | None:
+    """Safely get Mlflow logger from Trainer loggers.
+
+    Args:
+        trainer: Pytorch Lightning trainer.
+
+    Returns:
+        An mlflow logger if available, else None.
+    """
+    if isinstance(trainer.logger, MLFlowLogger):
+        return trainer.logger
+
+    if isinstance(trainer.logger, list):
+        for logger in trainer.logger:
+            if isinstance(logger, MLFlowLogger):
+                return logger
+
+    return None
diff --git a/quadra/utils/model_manager.py b/quadra/utils/model_manager.py
new file mode 100644
index 00000000..1dbfce1f
--- /dev/null
+++ b/quadra/utils/model_manager.py
@@ -0,0 +1,321 @@
+from __future__ import annotations
+
+import getpass
+import os
+from abc import ABC, abstractmethod
+from datetime import datetime
+from typing import Any, Literal
+
+from quadra.utils.utils import get_logger
+
+log = get_logger(__name__)
+
+try:
+    import mlflow  # noqa
+    from mlflow.entities import Run  # noqa
+    from mlflow.entities.model_registry import ModelVersion  # noqa
+    from mlflow.exceptions import RestException  # noqa
+    from mlflow.tracking import MlflowClient  # noqa
+
+    MLFLOW_AVAILABLE = True
+except ImportError:
+    MLFLOW_AVAILABLE = False
+
+
+VERSION_MD_TEMPLATE = "## **Version {}**\n"
+DESCRIPTION_MD_TEMPLATE = "### Description: \n{}\n"
+
+
+class AbstractModelManager(ABC):
+    """Abstract class for model managers."""
+
+    @abstractmethod
+    def register_model(
+        self, model_location: str, model_name: str, description: str, tags: dict[str, Any] | None = None
+    ) -> Any:
+        """Register a model in the model registry."""
+
+    @abstractmethod
+    def get_latest_version(self, model_name: str) -> Any:
+        """Get the latest version of a model for all the possible stages or filtered by stage."""
+
+    @abstractmethod
+    def transition_model(self, model_name: str, version: int, stage: str, description: str | None = None) -> Any:
+        """Transition the model with the given version to a new stage."""
+
+    @abstractmethod
+    def delete_model(self, model_name: str, version: int, description: str | None = None) -> None:
+        """Delete a model with the given version."""
+
+    @abstractmethod
+    def register_best_model(
+        self,
+        experiment_name: str,
+        metric: str,
+        model_name: str,
+        description: str,
+        tags: dict[str, Any] | None = None,
+        mode: Literal["max", "min"] = "max",
+        model_path: str = "deployment_model",
+    ) -> Any:
+        """Register the best model from an experiment."""
+
+    @abstractmethod
+    def download_model(self, model_name: str, version: int, output_path: str) -> None:
+        """Download the model with the given version to the given output path."""
+
+
+class MlflowModelManager(AbstractModelManager):
+    """Model manager for Mlflow."""
+
+    def __init__(self):
+        if not MLFLOW_AVAILABLE:
+            raise ImportError("Mlflow is not available, please install it with pip install mlflow")
+
+        if os.getenv("MLFLOW_TRACKING_URI") is None:
+            raise ValueError("MLFLOW_TRACKING_URI environment variable is not set")
+
+        mlflow.set_tracking_uri(os.getenv("MLFLOW_TRACKING_URI"))
+        self.client = MlflowClient()
+
+    def register_model(
+        self, model_location: str, model_name: str, description: str | None = None, tags: dict[str, Any] | None = None
+    ) -> ModelVersion:
+        """Register a model in the model registry.
+
+        Args:
+            model_location: The model uri
+            model_name: The name of the model after it is registered
+            description: A description of the model, this will be added to the model changelog
+            tags: A dictionary of tags to add to the model
+
+        Returns:
+            The model version
+        """
+        model_version = mlflow.register_model(model_uri=model_location, name=model_name, tags=tags)
+        log.info("Registered model %s with version %s", model_name, model_version.version)
+        registered_model_description = self.client.get_registered_model(model_name).description
+
+        if model_version.version == "1":
+            header = "# MODEL CHANGELOG\n"
+        else:
+            header = ""
+
+        new_model_description = VERSION_MD_TEMPLATE.format(model_version.version)
+        new_model_description += self._get_author_and_date()
+        new_model_description += self._generate_description(description)
+
+        self.client.update_registered_model(model_name, header + registered_model_description + new_model_description)
+
+        self.client.update_model_version(
+            model_name, model_version.version, "# MODEL CHANGELOG\n" + new_model_description
+        )
+
+        return model_version
+
+    def get_latest_version(self, model_name: str) -> ModelVersion:
+        """Get the latest version of a model.
+
+        Args:
+            model_name: The name of the model
+
+        Returns:
+            The model version
+        """
+        latest_version = max(int(x.version) for x in self.client.get_latest_versions("manager_model"))
+        model_version = self.client.get_model_version(model_name, latest_version)
+
+        return model_version
+
+    def transition_model(
+        self, model_name: str, version: int, stage: str, description: str | None = None
+    ) -> ModelVersion | None:
+        """Transition a model to a new stage.
+
+        Args:
+            model_name: The name of the model
+            version: The version of the model
+            stage: The stage of the model
+            description: A description of the transition, this will be added to the model changelog
+        """
+        previous_stage = self._safe_get_stage(model_name, version)
+
+        if previous_stage is None:
+            return None
+
+        if previous_stage.lower() == stage.lower():
+            log.warning("Model %s version %s is already in stage %s", model_name, version, stage)
+            return self.client.get_model_version(model_name, version)
+
+        log.info("Transitioning model %s version %s from %s to %s", model_name, version, previous_stage, stage)
+        model_version = self.client.transition_model_version_stage(name=model_name, version=version, stage=stage)
+        new_stage = model_version.current_stage
+        registered_model_description = self.client.get_registered_model(model_name).description
+        single_model_description = self.client.get_model_version(model_name, version).description
+
+        new_model_description = "## **Transition:**\n"
+        new_model_description += f"### Version {model_version.version} from {previous_stage} to {new_stage}\n"
+        new_model_description += self._get_author_and_date()
+        new_model_description += self._generate_description(description)
+
+        self.client.update_registered_model(model_name, registered_model_description + new_model_description)
+        self.client.update_model_version(
+            model_name, model_version.version, single_model_description + new_model_description
+        )
+
+        return model_version
+
+    def delete_model(self, model_name: str, version: int, description: str | None = None) -> None:
+        """Delete a model.
+
+        Args:
+            model_name: The name of the model
+            version: The version of the model
+            description: Why the model was deleted, this will be added to the model changelog
+        """
+        model_stage = self._safe_get_stage(model_name, version)
+
+        if model_stage is None:
+            return
+
+        if (
+            input(
+                f"Model named `{model_name}`, version {version} is in stage {model_stage}, "
+                "type the model name to continue deletion:"
+            )
+            != model_name
+        ):
+            log.warning("Model name did not match, aborting deletion")
+            return
+
+        log.info("Deleting model %s version %s", model_name, version)
+        self.client.delete_model_version(model_name, version)
+
+        registered_model_description = self.client.get_registered_model(model_name).description
+
+        new_model_description = "## **Deletion:**\n"
+        new_model_description += f"### Version {version} from stage: {model_stage}\n"
+        new_model_description += self._get_author_and_date()
+        new_model_description += self._generate_description(description)
+
+        self.client.update_registered_model(model_name, registered_model_description + new_model_description)
+
+    def register_best_model(
+        self,
+        experiment_name: str,
+        metric: str,
+        model_name: str,
+        description: str | None = None,
+        tags: dict[str, Any] | None = None,
+        mode: Literal["max", "min"] = "max",
+        model_path: str = "deployment_model",
+    ) -> ModelVersion | None:
+        """Register the best model from an experiment.
+
+        Args:
+            experiment_name: The name of the experiment
+            metric: The metric to use to determine the best model
+            model_name: The name of the model after it is registered
+            description: A description of the model, this will be added to the model changelog
+            tags: A dictionary of tags to add to the model
+            mode: The mode to use to determine the best model, either "max" or "min"
+            model_path: The path to the model within the experiment run
+
+        Returns:
+            The registered model version if successful, otherwise None
+        """
+        if mode not in ["max", "min"]:
+            raise ValueError(f"Mode must be either 'max' or 'min', got {mode}")
+
+        experiment_id = self.client.get_experiment_by_name(experiment_name).experiment_id
+        runs = self.client.search_runs(experiment_ids=[experiment_id])
+
+        if len(runs) == 0:
+            log.error("No runs found for experiment %s", experiment_name)
+            return None
+
+        best_run: Run | None = None
+
+        # We can only make comparisons if the model is on the top folder, otherwise just check if the folder exists
+        # TODO: Is there a better way to do this?
+        base_model_path = model_path.split("/")[0]
+
+        for run in runs:
+            run_artifacts = [x.path for x in self.client.list_artifacts(run.info.run_id) if x.path == base_model_path]
+
+            if len(run_artifacts) == 0:
+                # If we don't find the given model path, skip this run
+                continue
+
+            if best_run is None:
+                # If we find a run with the model it must also have the metric
+                if run.data.metrics.get(metric) is not None:
+                    best_run = run
+                continue
+
+            if mode == "max":
+                if run.data.metrics[metric] > best_run.data.metrics[metric]:
+                    best_run = run
+            else:
+                if run.data.metrics[metric] < best_run.data.metrics[metric]:
+                    best_run = run
+
+        if best_run is None:
+            log.error("No runs found for experiment %s with the given metric", experiment_name)
+            return None
+
+        best_model_uri = f"runs:/{best_run.info.run_id}/{model_path}"
+
+        model_version = self.register_model(
+            model_location=best_model_uri, model_name=model_name, tags=tags, description=description
+        )
+
+        return model_version
+
+    def download_model(self, model_name: str, version: int, output_path: str) -> None:
+        """Download the model with the given version to the given output path.
+
+        Args:
+            model_name: The name of the model
+            version: The version of the model
+            output_path: The path to save the model to
+        """
+        artifact_uri = self.client.get_model_version_download_uri(model_name, version)
+        log.info("Downloading model %s version %s from %s to %s", model_name, version, artifact_uri, output_path)
+        if not os.path.exists(output_path):
+            log.info("Creating output path %s", output_path)
+            os.makedirs(output_path)
+        mlflow.artifacts.download_artifacts(artifact_uri=artifact_uri, dst_path=output_path)
+
+    @staticmethod
+    def _generate_description(description: str | None = None) -> str:
+        """Generate the description markdown template."""
+        if description is None:
+            return ""
+
+        return DESCRIPTION_MD_TEMPLATE.format(description)
+
+    @staticmethod
+    def _get_author_and_date() -> str:
+        """Get the author and date markdown template."""
+        author_and_date = f"### Author: {getpass.getuser()}\n"
+        author_and_date += f"### Date: {datetime.now().astimezone().strftime('%d/%m/%Y %H:%M:%S %Z')}\n"
+
+        return author_and_date
+
+    def _safe_get_stage(self, model_name: str, version: int) -> str | None:
+        """Get the stage of a model version.
+
+        Args:
+            model_name: The name of the model
+            version: The version of the model
+
+        Returns:
+            The stage of the model version if it exists, otherwise None
+        """
+        try:
+            model_stage = self.client.get_model_version(model_name, version).current_stage
+            return model_stage
+        except RestException:
+            log.error("Model named %s with version %s does not exist", model_name, version)
+            return None
diff --git a/quadra/utils/models.py b/quadra/utils/models.py
index 0ffbd140..baa5dc78 100644
--- a/quadra/utils/models.py
+++ b/quadra/utils/models.py
@@ -1,15 +1,22 @@
 import math
 import warnings
-from typing import Callable, List, Tuple, Type, Union, cast
+from typing import Callable, List, Optional, Tuple, Type, Union, cast
 
 import numpy as np
+import timm
 import torch
 import torch.nn.functional as F
 import tqdm
+from pytorch_grad_cam import GradCAM
+from scipy import ndimage
+from sklearn.linear_model._base import ClassifierMixin
 from timm.models.layers import DropPath
 from timm.models.vision_transformer import Mlp
 from torch import nn
 
+from quadra.utils import utils
+from quadra.utils.vit_explainability import VitAttentionGradRollout
+
 
 def net_hat(input_size: int, output_size: int) -> torch.nn.Sequential:
     """Create a linear layer with input and output neurons.
@@ -70,8 +77,13 @@ def init_weights(m):
 
 
 def get_feature(
-    feature_extractor: torch.nn.Module, dl: torch.utils.data.DataLoader, iteration_over_training: int = 1
-) -> Tuple[np.ndarray, np.ndarray]:
+    feature_extractor: torch.nn.Module,
+    dl: torch.utils.data.DataLoader,
+    iteration_over_training: int = 1,
+    gradcam: bool = False,
+    classifier: Optional[ClassifierMixin] = None,
+    input_shape: Optional[Tuple[int, int, int]] = None,
+) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
     """Given a dataloader and a PyTorch model, extract features with the model and return features and labels.
 
     Args:
@@ -79,35 +91,96 @@ def get_feature(
         feature_extractor: Pretrained PyTorch backbone
         iteration_over_training: Extract feature iteration_over_training times for each image
             (best if used with augmentation)
+        gradcam: Whether to compute gradcams. Notice that it will slow the function
+        classifier: Scikit-learn classifier
+        input_shape: [H,W,C], backbone input shape, needed by classifier's pytorch wrapper
 
     Returns:
         Tuple containing:
             features: Model features
             labels: input_labels
+            grayscale_cams: Gradcam output maps, None if gradcam arg is False
     """
     feature_extractor.eval()
 
+    # Setup gradcam
+    if gradcam:
+        if isinstance(feature_extractor.features_extractor, timm.models.resnet.ResNet):
+            target_layers = [feature_extractor.features_extractor.layer4[-1]]
+            cam = GradCAM(
+                model=feature_extractor,
+                target_layers=target_layers,
+                use_cuda=torch.cuda.is_available(),
+            )
+            for p in feature_extractor.features_extractor.layer4[-1].parameters():
+                p.requires_grad = True
+        elif is_vision_transformer(feature_extractor.features_extractor):  # type: ignore[arg-type]
+            grad_rollout = VitAttentionGradRollout(
+                feature_extractor.features_extractor,  # type: ignore[arg-type]
+                classifier=classifier,
+                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
+
     # Extract features from data
-    with torch.no_grad():
-        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)
-                y_hat = cast(Union[List[torch.Tensor], Tuple[torch.Tensor], torch.Tensor], feature_extractor(x1))
-                if isinstance(y_hat, (list, tuple)):
-                    y_hat = y_hat[0].cpu()
-                else:
-                    y_hat = y_hat.cpu()
 
-                if i == 0 and iteration == 0:
-                    features = torch.cat([y_hat], dim=0)
-                    labels = np.concatenate([y1])
+    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 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]
+                    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
+                    orig_shape = grayscale_cam_low_res.shape
+                    new_shape = (orig_shape[0], x1.shape[2], x1.shape[3])
+                    zoom_factors = tuple(np.array(new_shape) / np.array(orig_shape))
+                    grayscale_cam = ndimage.zoom(grayscale_cam_low_res, zoom_factors, order=1)
                 else:
-                    features = torch.cat([features, y_hat], dim=0)
-                    labels = np.concatenate([labels, y1], axis=0)
+                    grayscale_cam = cam(input_tensor=x1, targets=None)
+                feature_extractor.zero_grad(set_to_none=True)
+                torch.cuda.empty_cache()
+
+            else:
+                with torch.no_grad():
+                    y_hat = cast(Union[List[torch.Tensor], Tuple[torch.Tensor], torch.Tensor], feature_extractor(x1))
+                grayscale_cams = None
+
+            if isinstance(y_hat, (list, tuple)):
+                y_hat = y_hat[0].cpu()
+            else:
+                y_hat = y_hat.cpu()
+
+            if i == 0 and iteration == 0:
+                features = torch.cat([y_hat], dim=0)
+                labels = np.concatenate([y1])
+                if gradcam:
+                    grayscale_cams = grayscale_cam
+            else:
+                features = torch.cat([features, y_hat], dim=0)
+                labels = np.concatenate([labels, y1], axis=0)
+                if gradcam:
+                    grayscale_cams = np.concatenate([grayscale_cams, grayscale_cam], axis=0)
+
+    return features.detach().numpy(), labels, grayscale_cams
+
+
+def is_vision_transformer(model: torch.nn.Module) -> bool:
+    """Verify if pytorch module is a Vision Transformer.
+    This check is primarily needed for gradcam computation in classification tasks.
 
-    return features.numpy(), labels
+    Args:
+        model: Model
+    """
+    return type(model).__name__ == "VisionTransformer"
 
 
 def _no_grad_trunc_normal_(tensor: torch.Tensor, mean: float, std: float, a: float, b: float):
@@ -274,7 +347,6 @@ class PositionalEncoding1D(torch.nn.Module):
     """
 
     def __init__(self, d_model: int, temperature: float = 10000.0, dropout: float = 0.0, max_len: int = 5000):
-
         super().__init__()
         self.dropout: Union[torch.nn.Dropout, torch.nn.Identity]
         if dropout > 0:
@@ -334,7 +406,6 @@ def __init__(
         mask_diagonal: bool = True,
         learnable_temperature: bool = True,
     ):
-
         super().__init__()
         self.norm1 = norm_layer(dim)
         self.attn = LocalSelfAttention(
diff --git a/quadra/utils/patch/dataset.py b/quadra/utils/patch/dataset.py
index 82f03b39..3453b645 100644
--- a/quadra/utils/patch/dataset.py
+++ b/quadra/utils/patch/dataset.py
@@ -371,7 +371,6 @@ def __save_patch_dataset(
                             discard_in_training = False
 
                 if save_mask:
-
                     mask_name = f"{image_name}_{row_index * image_patches.shape[1] + col_index}{mask_extension}.png"
 
                     if discard_in_training:
@@ -1268,7 +1267,8 @@ def load_train_file(
     exclude_filter: Optional[List[str]] = None,
     class_to_skip: Optional[list] = None,
 ) -> Tuple[List[str], List[str]]:
-    """Load a train file and return a list of samples and a list of targets.
+    """Load a train file and return a list of samples and a list of targets. It is expected that train files will be in
+        the same location as the train_file_path.
 
     Args:
         train_file_path: Training file location
@@ -1305,6 +1305,9 @@ def load_train_file(
     samples = [samples[i] for i in valid_samples_indices]
     targets = [targets[i] for i in valid_samples_indices]
 
+    train_folder = os.path.dirname(train_file_path)
+    samples = [os.path.join(train_folder, x) for x in samples]
+
     return samples, targets
 
 
diff --git a/quadra/utils/utils.py b/quadra/utils/utils.py
index 0849d409..5fda5928 100644
--- a/quadra/utils/utils.py
+++ b/quadra/utils/utils.py
@@ -1,16 +1,18 @@
 """Common utility functions.
 Some of them are mostly based on https://github.com/ashleve/lightning-hydra-template.
 """
+import glob
 import json
 import logging
 import os
 import subprocess
 import sys
 import warnings
-from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, cast
 
 import cv2
 import dotenv
+import mlflow
 import numpy as np
 import pytorch_lightning as pl
 import rich.syntax
@@ -23,7 +25,11 @@
 from pytorch_lightning.utilities import rank_zero_only
 
 import quadra
-from quadra.callbacks.mlflow import check_file_server_dependencies, check_minio_credentials, get_mlflow_logger
+import quadra.utils.export as quadra_export
+from quadra.callbacks.mlflow import get_mlflow_logger
+from quadra.utils.mlflow import infer_signature_torch_model
+
+IMAGE_EXTENSIONS: List[str] = [".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".pbm", ".pgm", ".ppm", ".pxm", ".pnm"]
 
 
 def get_logger(name=__name__) -> logging.Logger:
@@ -73,13 +79,6 @@ def extras(config: DictConfig) -> None:
         if config.datamodule.get("num_workers"):
             config.datamodule.num_workers = 0
 
-    if config.core.get("upload_artifacts"):
-        if config.logger.get("mlflow") is not None and config.logger.mlflow.get("_target_") is not None:
-            if "http" in config.logger.mlflow.get("tracking_uri"):
-                log.info("Checking mlfow settings for artifact upload since <config.core.upload_artifacts=True>")
-                check_file_server_dependencies()
-                check_minio_credentials()
-
 
 @rank_zero_only
 def print_config(
@@ -222,44 +221,78 @@ def upload_file_tensorboard(file_path: str, tensorboard_logger: TensorBoardLogge
 
 def finish(
     config: DictConfig,
-    model: pl.LightningModule,
+    module: pl.LightningModule,
     datamodule: pl.LightningDataModule,
     trainer: pl.Trainer,
     callbacks: List[pl.Callback],
     logger: List[pl.loggers.Logger],
+    export_folder: str,
 ) -> None:
     """Upload config files to MLFlow server.
 
     Args:
         config: Configuration composed by Hydra.
-        model: LightningModule.
+        module: LightningModule.
         datamodule: LightningDataModule.
         trainer: LightningTrainer.
         callbacks: List of LightningCallbacks.
         logger: List of LightningLoggers.
+        export_folder: Folder where the deployment models are exported.
     """
     # pylint: disable=unused-argument
+
     if len(logger) > 0 and config.core.get("upload_artifacts"):
-        mflow_logger = get_mlflow_logger(trainer=trainer)
+        mlflow_logger = get_mlflow_logger(trainer=trainer)
         tensorboard_logger = get_tensorboard_logger(trainer=trainer)
+        file_names = ["config.yaml", "config_resolved.yaml", "config_tree.txt", "data/dataset.csv"]
 
-        if mflow_logger is not None:
+        if mlflow_logger is not None:
             config_paths = []
-            file_names = ["config.yaml", "config_resolved.yaml", "config_tree.txt"]
+
             for f in file_names:
                 if os.path.isfile(os.path.join(os.getcwd(), f)):
                     config_paths.append(os.path.join(os.getcwd(), f))
+
             for path in config_paths:
-                mflow_logger.experiment.log_artifact(
-                    run_id=mflow_logger.run_id, local_path=path, artifact_path="metadata"
+                mlflow_logger.experiment.log_artifact(
+                    run_id=mlflow_logger.run_id, local_path=path, artifact_path="metadata"
                 )
 
+            deployed_models = glob.glob(os.path.join(export_folder, "*"))
+            model_json: Optional[Dict[str, Any]] = None
+
+            if os.path.exists(os.path.join(export_folder, "model.json")):
+                with open(os.path.join(export_folder, "model.json"), "r") as json_file:
+                    model_json = json.load(json_file)
+
+            if model_json is not None:
+                for model_path in deployed_models:
+                    if model_path.endswith(".pt"):
+                        model, _ = quadra_export.import_deployment_model(model_path, device="cpu")
+
+                        input_size = model_json["input_size"]
+
+                        # Not a huge fan of this check
+                        if not isinstance(input_size[0], list):
+                            # Input size is not a list of lists
+                            input_size = [input_size]
+
+                        inputs = cast(List[Any], quadra_export.generate_torch_inputs(input_size, device="cpu"))
+                        signature = infer_signature_torch_model(model, inputs)
+
+                        with mlflow.start_run(run_id=mlflow_logger.run_id) as _:
+                            mlflow.pytorch.log_model(
+                                model,
+                                artifact_path=model_path,
+                                signature=signature,
+                            )
+
         if tensorboard_logger is not None:
             config_paths = []
-            file_names = ["config.yaml", "config_resolved.yaml", "config_tree.txt"]
             for f in file_names:
                 if os.path.isfile(os.path.join(os.getcwd(), f)):
                     config_paths.append(os.path.join(os.getcwd(), f))
+
             for path in config_paths:
                 upload_file_tensorboard(file_path=path, tensorboard_logger=tensorboard_logger)
 
diff --git a/quadra/utils/vit_explainability.py b/quadra/utils/vit_explainability.py
new file mode 100644
index 00000000..bfaeb0b7
--- /dev/null
+++ b/quadra/utils/vit_explainability.py
@@ -0,0 +1,337 @@
+# Original Author: jacobgil
+# Title: Explainability for Vision Transformers
+# Source: https://github.com/jacobgil/vit-explain (MIT license)
+# Description: This is a heavily modified version of the original jacobgil code (the underlying math is still the same).
+
+import math
+from typing import List, Optional
+
+import numpy as np
+import torch
+from sklearn.linear_model._base import LinearClassifierMixin
+
+
+def rollout(
+    attentions: List[torch.Tensor], discard_ratio: float = 0.9, head_fusion: str = "mean", aspect_ratio: float = 1.0
+) -> np.ndarray:
+    """Apply rollout on Attention matrices.
+
+    Args:
+        attentions: List of Attention matrices coming from different blocks
+        discard_ratio: Percentage of elements to discard
+        head_fusion: Strategy of fusion of attention heads
+        aspect_ratio: Model inputs' width divided by height
+
+    Returns:
+        mask: Output mask, still needs a resize
+    """
+    result = torch.eye(attentions[0].size(-1))
+    with torch.no_grad():
+        for attention in attentions:
+            if head_fusion == "mean":
+                attention_heads_fused = attention.mean(dim=1)
+            elif head_fusion == "max":
+                attention_heads_fused = attention.max(dim=1)[0]
+            elif head_fusion == "min":
+                attention_heads_fused = attention.min(dim=1)[0]
+            else:
+                raise ValueError("Attention head fusion type Not supported")
+            # Drop the lowest attentions, but
+            # don't drop the class token
+            flat = attention_heads_fused.view(attention_heads_fused.size(0), -1)
+            _, indices = flat.topk(int(flat.size(-1) * discard_ratio), -1, False)
+            flat.scatter_(-1, indices, 0)
+            I = torch.eye(attention_heads_fused.size(-1))
+            a = (attention_heads_fused + 1.0 * I) / 2
+            a = a / a.sum(dim=-1).unsqueeze(1)
+            result = torch.matmul(a, result)
+    # Look at the total attention between the class token and the image patches
+    mask = result[:, 0, 1:]
+    batch_size = mask.size(0)
+    # TODO: Non squared input-size handling can be improved. Not easy though
+    height = math.floor((mask.size(-1) / aspect_ratio) ** 0.5)
+    total_size = mask.size(-1)
+    width = math.floor(total_size / height)
+    if mask.size(-1) > (height * width):
+        to_remove = mask.size(-1) - (height * width)
+        mask = mask[:, :-to_remove]
+    mask = mask.reshape(batch_size, height, width).numpy()
+    mask = mask / mask.max(axis=(1, 2), keepdims=True)
+
+    return mask
+
+
+class VitAttentionRollout:
+    """Attention gradient rollout class. Constructor registers hooks to the model's specified layers.
+    Only 4 layers by default given the high load on gpu. Best gradcams obtained using all blocks.
+
+    Args:
+        model: Model
+        attention_layer_names: On which layers to register the hook
+        head_fusion: Strategy of fusion for attention heads
+        discard_ratio: Percentage of elements to discard
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        attention_layer_names: Optional[List[str]] = None,
+        head_fusion: str = "mean",
+        discard_ratio: float = 0.9,
+    ):
+        if attention_layer_names is None:
+            attention_layer_names = [
+                "blocks.6.attn.attn_drop",
+                "blocks.7.attn.attn_drop",
+                "blocks.10.attn.attn_drop",
+                "blocks.11.attn.attn_drop",
+            ]
+        self.model = model
+        self.head_fusion = head_fusion
+        self.discard_ratio = discard_ratio
+        self.f_hook_handles: List[torch.utils.hooks.RemovableHandle] = []
+        for name, module in self.model.named_modules():
+            for layer_name in attention_layer_names:
+                if layer_name in name:
+                    self.f_hook_handles.append(module.register_forward_hook(self.get_attention))
+        self.attentions: List[torch.Tensor] = []
+
+    def get_attention(
+        self, module: torch.nn.Module, inpt: torch.Tensor, out: torch.Tensor  # pylint: disable=W0613
+    ) -> None:
+        """Hook to return attention.
+
+        Args:
+            module: Torch module
+            inpt: Input tensor
+            out: Output tensor, in this case the attention
+        """
+        self.attentions.append(out.detach().clone().cpu())
+
+    def __call__(self, input_tensor: torch.Tensor) -> np.ndarray:
+        """Called when the class instance is used as a function.
+
+        Args:
+            input_tensor: Input tensor
+
+        Returns:
+            out: Batch of output masks
+        """
+        self.attentions.clear()
+        with torch.no_grad():
+            self.model(input_tensor)
+        out = rollout(
+            self.attentions,
+            self.discard_ratio,
+            self.head_fusion,
+            aspect_ratio=(input_tensor.shape[-1] / input_tensor.shape[-2]),
+        )
+
+        return out
+
+
+############################################### GRAD ROLLOUT ##########################################################
+
+
+def grad_rollout(
+    attentions: List[torch.Tensor], gradients: List[torch.Tensor], discard_ratio: float = 0.9, aspect_ratio: float = 1.0
+) -> np.ndarray:
+    """Apply gradient rollout on Attention matrices.
+
+    Args:
+        attentions: Attention matrices
+        gradients: Target class gradient matrices
+        discard_ratio: Percentage of elements to discard
+        aspect_ratio: Model inputs' width divided by height
+
+    Returns:
+        mask: Output mask, still needs a resize
+    """
+    result = torch.eye(attentions[0].size(-1))
+    with torch.no_grad():
+        for attention, grad in zip(attentions, gradients):
+            weights = grad
+            attention_heads_fused = torch.mean((attention * weights), dim=1)
+            attention_heads_fused[attention_heads_fused < 0] = 0
+            # Drop the lowest attentions, but
+            # don't drop the class token
+            flat = attention_heads_fused.view(attention_heads_fused.size(0), -1)
+            _, indices = flat.topk(int(flat.size(-1) * discard_ratio), -1, False)
+            flat.scatter_(-1, indices, 0)
+            I = torch.eye(attention_heads_fused.size(-1))
+            a = (attention_heads_fused + 1.0 * I) / 2
+            a = a / a.sum(dim=-1).unsqueeze(1)
+            result = torch.matmul(a, result)
+    # Look at the total attention between the class token,
+    # and the image patches
+    mask = result[:, 0, 1:]
+    batch_size = mask.size(0)
+    # TODO: Non squared input-size handling can be improved. Not easy though
+    height = math.floor((mask.size(-1) / aspect_ratio) ** 0.5)
+    total_size = mask.size(-1)
+    width = math.floor(total_size / height)
+    if mask.size(-1) > (height * width):
+        to_remove = mask.size(-1) - (height * width)
+        mask = mask[:, :-to_remove]
+    mask = mask.reshape(batch_size, height, width).numpy()
+    mask = mask / mask.max(axis=(1, 2), keepdims=True)
+
+    return mask
+
+
+class VitAttentionGradRollout:
+    """Attention gradient rollout class. Constructor registers hooks to the model's specified layers.
+    Only 4 layers by default given the high load on gpu. Best gradcams obtained using all blocks.
+
+    Args:
+        model: Pytorch model
+        attention_layer_names: On which layers to register the hooks
+        discard_ratio: Percentage of elements to discard
+        classifier: Scikit-learn classifier. Leave it to None if model already has a classifier on top.
+    """
+
+    def __init__(  # pylint: disable=W0102
+        self,
+        model: torch.nn.Module,
+        attention_layer_names: Optional[List[str]] = None,
+        discard_ratio: float = 0.9,
+        classifier: Optional[LinearClassifierMixin] = None,
+        example_input: Optional[torch.Tensor] = None,
+    ):
+        if attention_layer_names is None:
+            attention_layer_names = [
+                "blocks.6.attn.attn_drop",
+                "blocks.7.attn.attn_drop",
+                "blocks.10.attn.attn_drop",
+                "blocks.11.attn.attn_drop",
+            ]
+
+        if classifier is not None:
+            if example_input is None:
+                raise ValueError(
+                    "Must provide an input example to LinearModelPytorchWrapper when classifier is not None"
+                )
+            self.model = LinearModelPytorchWrapper(
+                backbone=model,
+                linear_classifier=classifier,
+                example_input=example_input,
+                device=next(model.parameters()).device,
+            )
+        else:
+            self.model = model  # type: ignore[assignment]
+
+        self.discard_ratio = discard_ratio
+        self.f_hook_handles: List[torch.utils.hooks.RemovableHandle] = []
+        self.b_hook_handles: List[torch.utils.hooks.RemovableHandle] = []
+        for name, module in self.model.named_modules():
+            for layer_name in attention_layer_names:
+                if layer_name in name:
+                    self.f_hook_handles.append(module.register_forward_hook(self.get_attention))
+                    self.b_hook_handles.append(module.register_backward_hook(self.get_attention_gradient))
+        self.attentions: List[torch.Tensor] = []
+        self.attention_gradients: List[torch.Tensor] = []
+        # Activate gradients
+        blocks_list = [x.split("blocks")[1].split(".attn")[0] for x in attention_layer_names]
+        for name, module in model.named_modules():
+            for p in module.parameters():
+                if "blocks" in name and any(x in name for x in blocks_list):
+                    p.requires_grad = True
+
+    def get_attention(
+        self, module: torch.nn.Module, inpt: torch.Tensor, out: torch.Tensor  # pylint: disable=W0613
+    ) -> None:
+        """Hook to return attention.
+
+        Args:
+            module: Torch module
+            inpt: Input tensor
+            out: Output tensor, in this case the attention
+        """
+        self.attentions.append(out.detach().clone().cpu())
+
+    def get_attention_gradient(
+        self, module: torch.nn.Module, grad_input: torch.Tensor, grad_output: torch.Tensor  # pylint: disable=W0613
+    ) -> None:
+        """Hook to return attention.
+
+        Args:
+            module: Torch module
+            grad_input: Gradients' input tensor
+            grad_output: Gradients' output tensor, in this case the attention
+        """
+        self.attention_gradients.append(grad_input[0].detach().clone().cpu())
+
+    def __call__(self, input_tensor: torch.Tensor, targets_list: List[int]) -> np.ndarray:
+        """Called when the class instance is used as a function.
+
+        Args:
+            input_tensor: Model's input tensor
+            targets_list: List of targets. If None, argmax is used
+
+        Returns:
+            out: Batch of output masks
+        """
+        self.attentions.clear()
+        self.attention_gradients.clear()
+
+        self.model.zero_grad(set_to_none=True)
+        self.model.to(input_tensor.device)
+        output = self.model(input_tensor).cpu()
+
+        class_mask = torch.zeros(output.shape)
+        if targets_list is None:
+            targets_list = output.argmax(dim=1)
+        class_mask[torch.arange(output.shape[0]), targets_list] = 1
+        loss = (output * class_mask).sum()
+        loss.backward()
+        out = grad_rollout(
+            self.attentions,
+            self.attention_gradients,
+            self.discard_ratio,
+            aspect_ratio=(input_tensor.shape[-1] / input_tensor.shape[-2]),
+        )
+
+        return out
+
+
+class LinearModelPytorchWrapper(torch.nn.Module):
+    """Pytorch wrapper for scikit-learn linear models.
+
+    Args:
+        backbone: Backbone
+        linear_classifier: ScikitLearn linear classifier model
+        example_input: Input example needed to obtain output shape
+        device: The device to use. Defaults to "cpu"
+    """
+
+    def __init__(
+        self,
+        backbone: torch.nn.Module,
+        linear_classifier: LinearClassifierMixin,
+        example_input: torch.Tensor,
+        device: torch.device,
+    ):
+        super().__init__()
+        self.device = device
+        self.backbone = backbone.to(device)
+        if not isinstance(linear_classifier, LinearClassifierMixin):
+            raise TypeError("Classifier is not of type LinearClassifierMixin.")
+        self.num_classes = len(linear_classifier.classes_)
+        self.linear_classifier = linear_classifier
+        with torch.no_grad():
+            output = self.backbone(example_input.to(device))
+            num_filters = output.shape[-1]
+
+        self.classifier = torch.nn.Linear(num_filters, self.num_classes).to(device)
+        self.classifier.weight.data = torch.from_numpy(linear_classifier.coef_).float()
+        self.classifier.bias.data = torch.from_numpy(linear_classifier.intercept_).float()
+
+    def forward(self, x):
+        if self.num_classes == 2:
+            class_one_probabilities = torch.nn.Sigmoid()(self.classifier(self.backbone(x)))
+            class_zero_probabilities = 1 - class_one_probabilities
+            two_class_probs = torch.cat((class_zero_probabilities, class_one_probabilities), dim=1)
+
+            return two_class_probs
+
+        return torch.nn.Softmax(dim=1)(self.classifier(self.backbone(x)))


Test patch:
diff --git a/quadra/configs/experiment/base/classification/sklearn_classification_patch_test.yaml b/quadra/configs/experiment/base/classification/sklearn_classification_patch_test.yaml
index e7283665..ba747fff 100644
--- a/quadra/configs/experiment/base/classification/sklearn_classification_patch_test.yaml
+++ b/quadra/configs/experiment/base/classification/sklearn_classification_patch_test.yaml
@@ -3,19 +3,16 @@
 defaults:
   - override /transforms: default_resize
   - override /task: sklearn_classification_patch_test
-  - override /backbone: dino_vitb8
   - override /datamodule: base/sklearn_classification_patch
   - override /trainer: sklearn_classification
 
-backbone:
-  model:
-    pretrained: true
-    freeze: true
-
 core:
   tag: "run"
   name: "sklearn-classification-patch-test"
 
+task:
+  model_path: ???
+
 datamodule:
   num_workers: 8
   batch_size: 32
diff --git a/quadra/configs/experiment/base/classification/sklearn_classification_test.yaml b/quadra/configs/experiment/base/classification/sklearn_classification_test.yaml
index 8a377431..b422ab14 100644
--- a/quadra/configs/experiment/base/classification/sklearn_classification_test.yaml
+++ b/quadra/configs/experiment/base/classification/sklearn_classification_test.yaml
@@ -3,19 +3,22 @@
 defaults:
   - override /transforms: default_resize
   - override /task: sklearn_classification_test
-  - override /backbone: dino_vitb8
   - override /trainer: sklearn_classification
   - override /datamodule: base/sklearn_classification
 
-backbone:
-  model:
-    pretrained: true
-    freeze: true
-
 core:
   tag: run
   name: sklearn-classification-test
 
+task:
+  device: cuda:0
+  gradcam: false
+  output:
+    folder: classification_experiment
+    report: true
+    example: true
+  model_path: ???
+
 datamodule:
   num_workers: 8
   batch_size: 32
diff --git a/quadra/configs/task/sklearn_classification_patch_test.yaml b/quadra/configs/task/sklearn_classification_patch_test.yaml
index 6831f382..08a62973 100644
--- a/quadra/configs/task/sklearn_classification_patch_test.yaml
+++ b/quadra/configs/task/sklearn_classification_patch_test.yaml
@@ -5,4 +5,4 @@ output:
   report: true
   example: true
   reconstruction_method: major_voting
-experiment_path:
+model_path: ???
diff --git a/quadra/configs/task/sklearn_classification_test.yaml b/quadra/configs/task/sklearn_classification_test.yaml
index 532a27f0..214fb51f 100644
--- a/quadra/configs/task/sklearn_classification_test.yaml
+++ b/quadra/configs/task/sklearn_classification_test.yaml
@@ -1,7 +1,8 @@
 _target_: quadra.tasks.classification.SklearnTestClassification
 device: cuda:0
+gradcam: false
 output:
   folder: classification_experiment
   report: true
   example: true
-experiment_path:
+model_path: ???
diff --git a/quadra/utils/tests/models.py b/quadra/utils/tests/models.py
new file mode 100644
index 00000000..bc62d494
--- /dev/null
+++ b/quadra/utils/tests/models.py
@@ -0,0 +1,27 @@
+from typing import Any
+
+import torch
+from torch import nn
+
+
+class SingleInputModel(nn.Module):
+    """Model taking a single input."""
+
+    def forward(self, x: Any):
+        return x
+
+
+class DoubleInputModel(nn.Module):
+    """Model taking two inputs."""
+
+    def forward(self, x: Any, y: Any):
+        return x, y
+
+
+class UnsupportedInputModel(nn.Module):
+    """Model taking an unsupported input."""
+
+    def forward(self, x: torch.Tensor, y: str):
+        y = f"unsupported input: {y}"
+
+        return x
diff --git a/tests/datamodules/test_classification.py b/tests/datamodules/test_classification.py
index b188d7b3..00cb6292 100644
--- a/tests/datamodules/test_classification.py
+++ b/tests/datamodules/test_classification.py
@@ -139,15 +139,22 @@ def test_classification_patch_datamodule(classification_patch_dataset: classific
         info = PatchDatasetInfo(**json.load(f))
 
     # train samples are named like imagename_class.h5
-    train_samples_df = datamodule.data[datamodule.data["split"] == "train"]["samples"].tolist()
-    datamodule_train_samples = set(["_".join(s.split("_")[0:-1]) for s in train_samples_df])
+    train_samples_df = datamodule.train_data["samples"].tolist()
+    datamodule_train_samples = set(
+        [os.path.splitext(os.path.basename("_".join(s.split("_")[0:-1])))[0] for s in train_samples_df]
+    )
     # val samples are named like imagename_patchnumber.xyz
-    val_samples_df = datamodule.data[datamodule.data["split"] == "val"]["samples"].tolist()
-    datamodule_val_samples = set(["_".join(os.path.basename(s).split("_")[0:-1]) for s in val_samples_df])
+    val_samples_df = datamodule.val_data["samples"].tolist()
+    datamodule_val_samples = set(
+        [os.path.splitext("_".join(os.path.basename(s).split("_")[0:-1]))[0] for s in val_samples_df]
+    )
     # test samples are named like imagename_patchnumber.xyz and may contain #DISCARD# in the name
-    test_samples_df = datamodule.data[datamodule.data["split"] == "test"]["samples"].tolist()
+    test_samples_df = datamodule.test_data["samples"].tolist()
     datamodule_test_samples = set(
-        ["_".join(os.path.basename(s).replace("#DISCARD#", "").split("_")[0:-1]) for s in test_samples_df]
+        [
+            os.path.splitext("_".join(os.path.basename(s).replace("#DISCARD#", "").split("_")[0:-1]))[0]
+            for s in test_samples_df
+        ]
     )
 
     train_filenames = set([os.path.splitext(os.path.basename(s.image_path))[0] for s in info.train_files])
diff --git a/tests/models/test_signature_wrapper.py b/tests/models/test_signature_wrapper.py
new file mode 100644
index 00000000..70ea015b
--- /dev/null
+++ b/tests/models/test_signature_wrapper.py
@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+import torch
+
+from quadra.models.base import ModelSignatureWrapper
+from quadra.utils.export import generate_torch_inputs
+from quadra.utils.tests.models import DoubleInputModel, SingleInputModel, UnsupportedInputModel
+
+
+@torch.inference_mode()
+def test_simple_model():
+    """Test the input shape retrieval for a simple model."""
+    model = ModelSignatureWrapper(SingleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    model(x)
+    assert model.input_shapes == [(3, 224, 224)]
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+    model(*inputs)
+
+
+@torch.inference_mode()
+def test_double_input_model():
+    """Test the input shape retrieval for a model with multiple inputs."""
+    model = ModelSignatureWrapper(DoubleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    model(x, y)
+    assert model.input_shapes == [(3, 224, 224), (3, 448, 448)]
+
+    model.input_shapes = None
+    model(y=y, x=x)
+    # Check that the order of the inputs does not matter
+    assert model.input_shapes == [(3, 224, 224), (3, 448, 448)]
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+    assert len(inputs) == 2
+    assert inputs[0].shape == (1, 3, 224, 224)
+    assert inputs[1].shape == (1, 3, 448, 448)
+    model(*inputs)
+
+
+@torch.inference_mode()
+def test_dict_input_model():
+    """Test the input shape retrieval for a model with a dict input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    model({"x": x})
+    assert model.input_shapes == [{"x": (3, 224, 224)}]
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+    assert len(inputs) == 1
+    assert isinstance(inputs[0], dict)
+    assert inputs[0]["x"].shape == (1, 3, 224, 224)
+    model(*inputs)
+
+
+@torch.inference_mode()
+def test_tuple_input_model():
+    """Test the input shape retrieval for a model with a tuple input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    model((x, y))
+    assert model.input_shapes == [((3, 224, 224), (3, 448, 448))]
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+    assert len(inputs) == 1
+    assert len(inputs[0]) == 2
+    assert isinstance(inputs[0], tuple)
+    assert inputs[0][0].shape == (1, 3, 224, 224)
+    assert inputs[0][1].shape == (1, 3, 448, 448)
+    model(*inputs)
+
+
+@torch.inference_mode()
+def test_list_of_tensors_input_model():
+    """Test the input shape retrieval for a model with a list of tensors input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    model([x, y])
+    assert model.input_shapes == [[(3, 224, 224), (3, 448, 448)]]
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+    assert len(inputs) == 1
+    assert isinstance(inputs[0], list)
+    assert len(inputs[0]) == 2
+    assert inputs[0][0].shape == (1, 3, 224, 224)
+    assert inputs[0][1].shape == (1, 3, 448, 448)
+    model(*inputs)
+
+
+@torch.inference_mode()
+def test_list_of_dicts_input_model():
+    """Test the input shape retrieval for a model with a list of dicts input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    model([{"x": x}, {"y": y}])
+    assert model.input_shapes == [[{"x": (3, 224, 224)}, {"y": (3, 448, 448)}]]
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+    assert len(inputs) == 1
+    assert isinstance(inputs[0], list)
+    assert len(inputs[0]) == 2
+    assert inputs[0][0]["x"].shape == (1, 3, 224, 224)
+    assert inputs[0][1]["y"].shape == (1, 3, 448, 448)
+    model(*inputs)
+
+
+@torch.inference_mode()
+def test_tuple_of_dicts_input_model():
+    """Test the input shape retrieval for a model with a tuple of dicts input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    model(({"x": x}, {"y": y}))
+    assert model.input_shapes == [({"x": (3, 224, 224)}, {"y": (3, 448, 448)})]
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+    assert len(inputs) == 1
+    assert isinstance(inputs[0], tuple)
+    assert len(inputs[0]) == 2
+    assert inputs[0][0]["x"].shape == (1, 3, 224, 224)
+    assert inputs[0][1]["y"].shape == (1, 3, 448, 448)
+    model(*inputs)
+
+
+@torch.inference_mode()
+def test_unsupported_input_model():
+    """Test the input shape retrieval for a model with an unsupported input."""
+    model = ModelSignatureWrapper(UnsupportedInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = "test"
+    model(x, y)
+    assert model.input_shapes is None
+    assert model.disable
diff --git a/tests/tasks/test_anomaly.py b/tests/tasks/test_anomaly.py
index 4a4d327b..27a44549 100644
--- a/tests/tasks/test_anomaly.py
+++ b/tests/tasks/test_anomaly.py
@@ -77,7 +77,6 @@ def test_padim_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_datas
         f"datamodule.data_path={data_path}",
         "model.model.backbone=resnet18",
         f"model.dataset.task={task}",
-        "task.export_type=[torchscript]",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
@@ -105,7 +104,6 @@ def test_patchcore_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_d
         f"datamodule.data_path={data_path}",
         "model.model.backbone=resnet18",
         f"model.dataset.task={task}",
-        "task.export_type=[torchscript]",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
@@ -131,7 +129,7 @@ def test_cflow_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_datas
         f"datamodule.data_path={data_path}",
         "model.model.backbone=resnet18",
         f"model.dataset.task={task}",
-        "task.export_type=null",
+        "task.export_config=null",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
@@ -157,7 +155,6 @@ def test_csflow_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_data
         "experiment=base/anomaly/csflow",
         f"datamodule.data_path={data_path}",
         f"model.dataset.task={task}",
-        "task.export_type=[torchscript]",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
@@ -185,7 +182,6 @@ def test_fastflow_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_da
         f"datamodule.data_path={data_path}",
         "model.model.backbone=resnet18",
         f"model.dataset.task={task}",
-        "task.export_type=[torchscript]",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
@@ -212,7 +208,6 @@ def test_draem_training(tmp_path: Path, base_anomaly_dataset: base_anomaly_datas
         "experiment=base/anomaly/draem",
         f"datamodule.data_path={data_path}",
         f"model.dataset.task={task}",
-        "task.export_type=[torchscript]",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
diff --git a/tests/tasks/test_classification.py b/tests/tasks/test_classification.py
index b27a2c4c..712e1e6f 100644
--- a/tests/tasks/test_classification.py
+++ b/tests/tasks/test_classification.py
@@ -50,12 +50,14 @@ 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.experiment_path={train_path}",
+        f"task.model_path={trained_model_path}",
         "backbone=resnet18",
         "task.device=cpu",
+        "task.gradcam=true",
     ] + BASE_EXPERIMENT_OVERRIDES
     execute_quadra_experiment(overrides=inference_overrides, experiment_path=test_path)
 
@@ -107,10 +109,11 @@ def test_inference_patches(tmp_path: Path, base_patch_classification_dataset: ba
     ] + 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.experiment_path={train_experiment_path}",
+        f"task.model_path={trained_model_path}",
         f"backbone={backbone}",
         "task.device=cpu",
     ] + BASE_EXPERIMENT_OVERRIDES
@@ -122,19 +125,20 @@ def test_inference_patches(tmp_path: Path, base_patch_classification_dataset: ba
 def _run_inference_experiment(data_path: str, train_path: str, test_path: str):
     test_overrides = [
         "experiment=base/classification/classification_evaluation",
-        "task.device=cpu",
         f"datamodule.data_path={data_path}",
         "datamodule.num_workers=1",
         "datamodule.batch_size=16",
         "logger=csv",
-        f"task.model_path={os.path.join(train_path, 'deployment_model', 'model.pt')}",
+        "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)
 
 
 @pytest.mark.parametrize(
-    "run_test, backbone, gradcam", [(True, "resnet18", True), (False, "resnet18", False), (True, "dino_vits8", False)]
+    "run_test, backbone, gradcam, freeze",
+    [(True, "resnet18", True, False), (False, "resnet18", False, False), (True, "dino_vits8", False, True)],
 )
 def test_train_classification(
     tmp_path: Path,
@@ -142,6 +146,7 @@ def test_train_classification(
     run_test: bool,
     backbone: str,
     gradcam: bool,
+    freeze: bool,
 ):
     data_path, arguments = base_classification_dataset
 
@@ -156,7 +161,8 @@ def test_train_classification(
         f"datamodule.data_path={data_path}",
         f"model.num_classes={num_classes}",
         f"backbone={backbone}",
-        f"model.module.gradcam={gradcam}",
+        f"backbone.model.freeze={freeze}",
+        f"task.gradcam={gradcam}",
         "trainer.max_epochs=1",
         "task.report=True",
         f"task.run_test={run_test}",
diff --git a/tests/tasks/test_segmentation.py b/tests/tasks/test_segmentation.py
index 38d207c8..75cb393d 100644
--- a/tests/tasks/test_segmentation.py
+++ b/tests/tasks/test_segmentation.py
@@ -1,4 +1,5 @@
 # pylint: disable=redefined-outer-name
+import os
 import shutil
 from pathlib import Path
 
@@ -17,7 +18,6 @@
     "+trainer.limit_val_batches=1",
     "+trainer.limit_test_batches=1",
     "logger=csv",
-    "task.evaluate.analysis=false",
 ]
 
 
@@ -27,14 +27,29 @@ def test_smp_binary(
 ):
     data_path, _, _ = base_binary_segmentation_dataset
 
+    train_path = tmp_path / "train"
+    test_path = tmp_path / "test"
+    train_path.mkdir()
+    test_path.mkdir()
+
     overrides = [
         "experiment=base/segmentation/smp",
         f"datamodule.data_path={data_path}",
         f"task.report={generate_report}",
+        "task.evaluate.analysis=false",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
-    execute_quadra_experiment(overrides=overrides, experiment_path=tmp_path)
+    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)
 
     shutil.rmtree(tmp_path)
 
@@ -43,6 +58,11 @@ def test_smp_multiclass(tmp_path: Path, base_multiclass_segmentation_dataset: ba
     data_path, _, class_to_idx = base_multiclass_segmentation_dataset
     idx_to_class = {v: k for k, v in class_to_idx.items()}
 
+    train_path = tmp_path / "train"
+    test_path = tmp_path / "test"
+    train_path.mkdir()
+    test_path.mkdir()
+
     idx_to_class_parameter = str(idx_to_class).replace(
         "'", ""
     )  # Remove single quotes so that it can be parsed by hydra
@@ -51,10 +71,21 @@ def test_smp_multiclass(tmp_path: Path, base_multiclass_segmentation_dataset: ba
         "experiment=base/segmentation/smp_multiclass",
         f"datamodule.data_path={data_path}",
         f"datamodule.idx_to_class={idx_to_class_parameter}",
+        "task.evaluate.analysis=false",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
-    execute_quadra_experiment(overrides=overrides, experiment_path=tmp_path)
+    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)
 
     shutil.rmtree(tmp_path)
 
@@ -73,6 +104,7 @@ def test_smp_multiclass_with_binary_dataset(
         "experiment=base/segmentation/smp_multiclass",
         f"datamodule.data_path={data_path}",
         f"datamodule.idx_to_class={idx_to_class_parameter}",
+        "task.evaluate.analysis=false",
     ]
     overrides += BASE_EXPERIMENT_OVERRIDES
 
diff --git a/tests/utilities/test_mlflow_export.py b/tests/utilities/test_mlflow_export.py
new file mode 100644
index 00000000..4a875c42
--- /dev/null
+++ b/tests/utilities/test_mlflow_export.py
@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import importlib
+
+import pytest
+
+try:
+    importlib.import_module("mlflow")
+except ImportError:
+    pytest.skip("Mlflow is not installed", allow_module_level=True)
+
+from typing import Sequence
+
+import torch
+from mlflow.models import infer_signature
+from mlflow.types.schema import TensorSpec
+from pytest import raises
+
+from quadra.models.base import ModelSignatureWrapper
+from quadra.utils.export import generate_torch_inputs
+from quadra.utils.mlflow import infer_signature_input_torch, infer_signature_torch_model
+from quadra.utils.tests.models import DoubleInputModel, SingleInputModel
+
+
+def check_signature_equality(
+    input_signature: Sequence[TensorSpec] | dict[str, TensorSpec] | TensorSpec,
+    expected_signature: Sequence[TensorSpec] | dict[str, TensorSpec] | TensorSpec,
+):
+    """Assert that the generated signature matches the expected signature."""
+    is_equal = True
+
+    if isinstance(input_signature, Sequence) and isinstance(expected_signature, Sequence):
+        for c_input_signature, c_expected_signature in zip(input_signature, expected_signature):
+            is_equal = is_equal and check_signature_equality(c_input_signature, c_expected_signature)
+    elif isinstance(input_signature, dict) and isinstance(expected_signature, dict):
+        for k, v in input_signature.items():
+            is_equal = is_equal and check_signature_equality(v, expected_signature[k])
+    elif isinstance(input_signature, TensorSpec) and isinstance(expected_signature, TensorSpec):
+        is_equal = (
+            input_signature.type == expected_signature.type
+            and input_signature.shape == expected_signature.shape
+            and input_signature.name == expected_signature.name
+        )
+    else:
+        is_equal = False
+
+    return is_equal
+
+
+@torch.inference_mode()
+def test_single_tensor_signature():
+    """Test the input shape retrieval for a simple model."""
+    model = ModelSignatureWrapper(SingleInputModel())
+
+    x = torch.zeros(1, 3, 224, 224)
+    _ = model(x)
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+
+    signature = infer_signature_torch_model(model, inputs)
+
+    expected_input_signature = [TensorSpec(shape=(-1, *x.shape[1:]), type=x.numpy().dtype)]
+
+    # The model output is the same as the input
+    assert check_signature_equality(signature.inputs.inputs, expected_input_signature)
+    assert check_signature_equality(signature.outputs.inputs, expected_input_signature)
+
+
+@torch.inference_mode()
+def test_multiple_tensor_signature():
+    """Test the input shape retrieval for a model with multiple inputs."""
+    model = ModelSignatureWrapper(DoubleInputModel())
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    _ = model(x, y)
+    assert model.input_shapes == [(3, 224, 224), (3, 448, 448)]
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+
+    signature = infer_signature_torch_model(model, inputs)
+
+    expected_input_signature = [
+        TensorSpec(shape=(-1, *x.shape[1:]), type=x.numpy().dtype, name="output_0"),
+        TensorSpec(shape=(-1, *y.shape[1:]), type=y.numpy().dtype, name="output_1"),
+    ]
+
+    # The model output is the same as the input
+    assert check_signature_equality(signature.inputs.inputs, expected_input_signature)
+    assert check_signature_equality(signature.outputs.inputs, expected_input_signature)
+
+
+@torch.inference_mode()
+def test_dict_signature():
+    """Test the input shape retrieval for a model with a dict input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+
+    x = torch.zeros(1, 3, 224, 224)
+    _ = model({"x": x})
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+
+    signature = infer_signature_torch_model(model, inputs)
+
+    expected_input_signature = [TensorSpec(shape=(-1, *x.shape[1:]), type=x.numpy().dtype, name="x")]
+
+    # The model output is the same as the input
+    assert check_signature_equality(signature.inputs.inputs, expected_input_signature)
+    assert check_signature_equality(signature.outputs.inputs, expected_input_signature)
+
+
+@torch.inference_mode()
+def test_nested_tuple_signature():
+    """Test the input shape retrieval for a model with a tuple input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+    assert model.input_shapes is None
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    outputs = model((x, y))
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+
+    expected_output_signature = [
+        TensorSpec(shape=(-1, *x.shape[1:]), type=x.numpy().dtype, name="output_0"),
+        TensorSpec(shape=(-1, *y.shape[1:]), type=y.numpy().dtype, name="output_1"),
+    ]
+
+    # Nested structures are not supported
+    with raises(ValueError):
+        infer_signature_input_torch(inputs)
+
+    signature = infer_signature(infer_signature_input_torch(outputs), infer_signature_input_torch(outputs))
+
+    assert check_signature_equality(signature.outputs.inputs, expected_output_signature)
+
+
+@torch.inference_mode()
+def test_nested_list_signature():
+    """Test the input shape retrieval for a model with a list of tensors input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    outputs = model([x, y])
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+
+    expected_output_signature = [
+        TensorSpec(shape=(-1, *x.shape[1:]), type=x.numpy().dtype, name="output_0"),
+        TensorSpec(shape=(-1, *y.shape[1:]), type=y.numpy().dtype, name="output_1"),
+    ]
+
+    # Nested structures are not supported
+    with raises(ValueError):
+        infer_signature_input_torch(inputs)
+
+    signature = infer_signature(infer_signature_input_torch(outputs), infer_signature_input_torch(outputs))
+
+    assert check_signature_equality(signature.outputs.inputs, expected_output_signature)
+
+
+@torch.inference_mode()
+def test_nested_dicts_signature():
+    """Test the input shape retrieval for a model with a list of dicts input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    outputs = model([{"x": x}, {"y": y}])
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+
+    # Nested structures are not supported
+    with raises(ValueError):
+        infer_signature_input_torch(inputs)
+
+    with raises(ValueError):
+        infer_signature_input_torch(outputs)
+
+
+@torch.inference_mode()
+def test_tuple_of_dicts_signature():
+    """Test the input shape retrieval for a model with a tuple of dicts input."""
+    model = ModelSignatureWrapper(SingleInputModel())
+
+    x = torch.zeros(1, 3, 224, 224)
+    y = torch.zeros(1, 3, 448, 448)
+    outputs = model(({"x": x}, {"y": y}))
+
+    inputs = generate_torch_inputs(model.input_shapes, device="cpu")
+
+    # Nested structures are not supported
+    with raises(ValueError):
+        infer_signature_input_torch(inputs)
+
+    with raises(ValueError):
+        infer_signature_input_torch(outputs)