diff --git a/docling/models/code_formula_model.py b/docling/models/code_formula_model.py
index 1eb9e74d..5c23dcc8 100644
--- a/docling/models/code_formula_model.py
+++ b/docling/models/code_formula_model.py
@@ -234,7 +234,7 @@ class CodeFormulaModel(BaseItemAndImageEnrichmentModel):
         ----------
         doc : DoclingDocument
             The document being processed.
-        element_batch : Iterable[NodeItem]
+        element_batch : Iterable[ItemAndImageEnrichmentElement]
             A batch of elements to be processed.
 
         Returns