diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7e0feea9..0360767f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,12 @@
 # Changelog
 All notable changes to this project will be documented in this file.
 
+### [2.0.3]
+
+#### Fixed 
+
+- Fix anomaly visualizer callback not working properly after lightning upgrade
+
 ### [2.0.2]
 
 #### Fixed
diff --git a/pyproject.toml b/pyproject.toml
index 7d2d6396..47c7ecd7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "quadra"
-version = "2.0.2"
+version = "2.0.3"
 description = "Deep Learning experiment orchestration library"
 authors = [
   "Federico Belotti <federico.belotti@orobix.com>",
diff --git a/quadra/__init__.py b/quadra/__init__.py
index 5e9a1eb0..0525ae75 100644
--- a/quadra/__init__.py
+++ b/quadra/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "2.0.2"
+__version__ = "2.0.3"
 
 
 def get_version():
diff --git a/quadra/callbacks/anomalib.py b/quadra/callbacks/anomalib.py
index 1cba3926..ce9197a5 100644
--- a/quadra/callbacks/anomalib.py
+++ b/quadra/callbacks/anomalib.py
@@ -138,23 +138,23 @@ def _add_images(self, visualizer: Visualizer, filename: Path, output_label_folde
 
     def on_test_batch_end(
         self,
-        _trainer: pl.Trainer,
+        trainer: pl.Trainer,
         pl_module: AnomalyModule,
         outputs: Optional[STEP_OUTPUT],
-        _batch: Any,
-        _batch_idx: int,
-        _dataloader_idx: int,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
     ) -> None:
         """Log images at the end of every batch.
 
         Args:
-            _trainer: Pytorch lightning trainer object (unused).
+            trainer: Pytorch lightning trainer object (unused).
             pl_module: Lightning modules derived from BaseAnomalyLightning object as
                 currently only they support logging images.
             outputs: Outputs of the current test step.
-            _batch: Input batch of the current test step (unused).
-            _batch_idx: Index of the current test batch (unused).
-            _dataloader_idx: Index of the dataloader that yielded the current batch (unused).
+            batch: Input batch of the current test step (unused).
+            batch_idx: Index of the current test batch (unused).
+            dataloader_idx: Index of the dataloader that yielded the current batch (unused).
         """
         if self.disable:
             return