diff --git a/src/f3dasm/_src/optimization/optimizer.py b/src/f3dasm/_src/optimization/optimizer.py
index b2ad3fd9..b3dc2b2e 100644
--- a/src/f3dasm/_src/optimization/optimizer.py
+++ b/src/f3dasm/_src/optimization/optimizer.py
@@ -46,6 +46,10 @@ def dfdx(self, x: np.ndarray) -> np.ndarray:
 
 class ExperimentData(Protocol):
 
+    def __init__(self, domain: Domain, input_data: np.ndarray,
+                 output_data: np.ndarray):
+        ...
+
     @property
     def domain(self) -> Domain:
         ...
@@ -290,7 +294,6 @@ def _iterate(self, iterations: int, kwargs: Dict[str, Any],
             new_samples = self.update_step()
 
             # If new_samples is a tuple of input_data and output_data
-            # TODO: This doesn't work with a Protocol
             if isinstance(new_samples, tuple):
                 new_samples = type(self.data)(
                     domain=self.data.domain,