diff --git a/luxonis_train/core/utils/infer_utils.py b/luxonis_train/core/utils/infer_utils.py
index 630d2ccc..ded13b41 100644
--- a/luxonis_train/core/utils/infer_utils.py
+++ b/luxonis_train/core/utils/infer_utils.py
@@ -15,7 +15,7 @@ def render_visualizations(
     visualizations: dict[str, dict[str, Tensor]],
     save_dir: str | Path | None,
     show: bool = True,
-) -> None:
+) -> dict[str, list[np.ndarray]]:
     """Render or save visualizations."""
     save_dir = Path(save_dir) if save_dir is not None else None
     if save_dir is not None: