diff --git a/docling/datamodel/document.py b/docling/datamodel/document.py
index 788cccae..b8177730 100644
--- a/docling/datamodel/document.py
+++ b/docling/datamodel/document.py
@@ -11,6 +11,7 @@ from docling_core.types import FileInfoObject as DsFileInfoObject
 from docling_core.types import PageDimensions, PageReference, Prov, Ref
 from docling_core.types import Table as DsSchemaTable
 from docling_core.types import TableCell
+from docling_core.types.doc.base import Figure
 from pydantic import BaseModel
 from typing_extensions import deprecated
 
@@ -279,7 +280,7 @@ class ConvertedDocument(BaseModel):
                     ),
                 )
                 figures.append(
-                    BaseCell(
+                    Figure(
                         prov=[
                             Prov(
                                 bbox=target_bbox,
@@ -344,7 +345,7 @@ class ConvertedDocument(BaseModel):
             "subtitle-level-1",
             "paragraph",
             "caption",
-        ]
+        ],
     ):
         return self.output.export_to_markdown(
             delim=delim,
diff --git a/examples/batch_convert.py b/examples/batch_convert.py
index b88236b1..4491b364 100644
--- a/examples/batch_convert.py
+++ b/examples/batch_convert.py
@@ -37,11 +37,11 @@ def export_documents(
             # Export Markdown format:
             with (output_dir / f"{doc_filename}.md").open("w") as fp:
                 fp.write(conv_res.render_as_markdown())
-                
+
             # Export Document Tags format:
             with (output_dir / f"{doc_filename}.doctags").open("w") as fp:
-                fp.write(conv_res.render_as_doctags())                
-                
+                fp.write(conv_res.render_as_doctags())
+
         elif conv_res.status == ConversionStatus.PARTIAL_SUCCESS:
             _log.info(
                 f"Document {conv_res.input.file} was partially converted with the following errors:"