From d8e32b87245f60ef804f8e069c9c473adcb6a68e Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Thu, 6 Nov 2025 15:57:42 -0700
Subject: [PATCH] Ensure after creating a class that things are correct

---
 web/src/views/classification/ModelTrainingView.tsx | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/web/src/views/classification/ModelTrainingView.tsx b/web/src/views/classification/ModelTrainingView.tsx
index 036fe0e56..614d23c96 100644
--- a/web/src/views/classification/ModelTrainingView.tsx
+++ b/web/src/views/classification/ModelTrainingView.tsx
@@ -118,6 +118,11 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) {
 
   const [trainFilter, setTrainFilter] = useApiFilter<TrainFilter>();
 
+  const refreshAll = useCallback(() => {
+    refreshTrain();
+    refreshDataset();
+  }, [refreshTrain, refreshDataset]);
+
   // image multiselect
 
   const [selectedImages, setSelectedImages] = useState<string[]>([]);
@@ -409,7 +414,7 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) {
           trainImages={trainImages || []}
           trainFilter={trainFilter}
           selectedImages={selectedImages}
-          onRefresh={refreshTrain}
+          onRefresh={refreshAll}
           onClickImages={onClickImages}
           onDelete={onDelete}
         />