diff --git a/DashAI/front/src/pages/models/ModelsContent.jsx b/DashAI/front/src/pages/models/ModelsContent.jsx
index 9dcea249c..32bb83394 100644
--- a/DashAI/front/src/pages/models/ModelsContent.jsx
+++ b/DashAI/front/src/pages/models/ModelsContent.jsx
@@ -923,10 +923,10 @@ export default function ModelsContent() {
                 ) : step === 2 && selectedDatasetId ? (
                   <DatasetVisualization
                     dataset={datasets.find((d) => d.id === selectedDatasetId)}
-                    onSessionCreated={handleSessionCreated}
-                    onNewSession={handleNewSessionFromDataset}
-                    existingSessions={sessions}
-                    tasks={tasks}
+                    onItemCreated={handleSessionCreated}
+                    onNewItem={handleNewSessionFromDataset}
+                    existingItems={sessions}
+                    newItemButtonText="New Session"
                   />
                 ) : step === 0 ? (
                   <SelectOptionMenu