diff --git a/cuvis/Viewer.py b/cuvis/Viewer.py
index 8450b22..c92e122 100644
--- a/cuvis/Viewer.py
+++ b/cuvis/Viewer.py
@@ -51,7 +51,7 @@ def _create_view_data(self, new_handle: int) -> Union[dict[str, ImageData], Imag
         else:
             return view_array
 
-    def apply(self, mesu: Measurement) -> dict[str, ImageData]:
+    def apply(self, mesu: Measurement) -> Union[dict[str, ImageData], ImageData]:
         _ptr = cuvis_il.new_p_int()
         if cuvis_il.status_ok != cuvis_il.cuvis_viewer_apply(self._handle,
                                                              mesu._handle, _ptr):