diff --git a/meridian/analysis/optimizer.py b/meridian/analysis/optimizer.py
index c9e2be51a..59b632cde 100644
--- a/meridian/analysis/optimizer.py
+++ b/meridian/analysis/optimizer.py
@@ -1917,7 +1917,7 @@ def _validate_grid(
 
   def create_optimization_grid(
       self,
-      new_data: xr.Dataset | None = None,
+      new_data: analyzer_module.DataTensors | None = None,
       use_posterior: bool = True,
       selected_geos: Sequence[str] | None = None,
       # TODO: b/409550413 - Remove this argument.