diff --git a/spikeinterface_gui/backend_panel.py b/spikeinterface_gui/backend_panel.py
index 7c4da4b..19c8a2b 100644
--- a/spikeinterface_gui/backend_panel.py
+++ b/spikeinterface_gui/backend_panel.py
@@ -431,8 +431,8 @@ def set_external_curation(self, curation_data):
 
         curation_view = self.views["curation"]
         self.controller.set_curation_data(curation_data)
-        self.controller.current_curation_saved = True
         curation_view.notify_manual_curation_updated()
+        self.controller.current_curation_saved = True
         curation_view.refresh()
 
         # we also need to refresh the unit list view to update the unit visibility according to the new curation
diff --git a/spikeinterface_gui/controller.py b/spikeinterface_gui/controller.py
index 105783f..3ddea2d 100644
--- a/spikeinterface_gui/controller.py
+++ b/spikeinterface_gui/controller.py
@@ -900,7 +900,7 @@ def set_curation_data(self, curation_data):
             new_curation_data["label_definitions"] = default_label_definitions.copy()
 
         # validate the curation data
-        model = CurationModel(**new_curation_data)
+        model = Curation(**new_curation_data)
         self.curation_data = model.model_dump()
 
     def save_curation_in_analyzer(self):
@@ -1117,11 +1117,13 @@ def set_label_to_unit(self, unit_id, category, label):
         if label is None:
             self.remove_category_from_unit(unit_id, category)
             return
+        
+        label_types = self.curation_data['label_definitions'].keys()
 
         ix = self.find_unit_in_manual_labels(unit_id)
         if ix is not None:
             lbl = self.curation_data["manual_labels"][ix]
-            if "labels" in lbl and category in lbl["labels"]:
+            if "labels" in lbl and category in label_types:
                 # v2 format
                 lbl["labels"][category] = [label]
             elif category in lbl: