diff --git a/lib/pages/spaces_management/assign_tag_models/bloc/assign_tag_model_bloc.dart b/lib/pages/spaces_management/assign_tag_models/bloc/assign_tag_model_bloc.dart
index c8c4cea4..b5251868 100644
--- a/lib/pages/spaces_management/assign_tag_models/bloc/assign_tag_model_bloc.dart
+++ b/lib/pages/spaces_management/assign_tag_models/bloc/assign_tag_model_bloc.dart
@@ -123,7 +123,6 @@ class AssignTagModelBloc
 
   bool _validateTags(List<TagModel> tags) {
     if (tags.isEmpty) {
-      print("tags empty");
       return false;
     }
     final uniqueTags = tags.map((tag) => tag.tag?.trim() ?? '').toSet();
diff --git a/lib/pages/spaces_management/assign_tag_models/views/assign_tag_models_dialog.dart b/lib/pages/spaces_management/assign_tag_models/views/assign_tag_models_dialog.dart
index 82a721f0..a6b6b7f0 100644
--- a/lib/pages/spaces_management/assign_tag_models/views/assign_tag_models_dialog.dart
+++ b/lib/pages/spaces_management/assign_tag_models/views/assign_tag_models_dialog.dart
@@ -232,8 +232,6 @@ class AssignTagModelsDialog extends StatelessWidget {
                         builder: (buttonContext) => CancelButton(
                           label: 'Add New Device',
                           onPressed: () async {
-                            Navigator.of(context).pop();
-
                             for (var tag in state.tags) {
                               if (tag.location == null || subspaces == null) {
                                 continue;
@@ -255,13 +253,13 @@ class AssignTagModelsDialog extends StatelessWidget {
                                 state.tags.removeWhere(
                                     (t) => t.internalId == tag.internalId);
                               }
-
+                            }
+                            if (context.mounted) {
                               await showDialog<bool>(
                                 barrierDismissible: false,
-                                context:
-                                    Navigator.of(context, rootNavigator: true)
-                                        .context,
-                                builder: (context) => AddDeviceTypeModelWidget(
+                                context: context,
+                                builder: (dialogContext) =>
+                                    AddDeviceTypeModelWidget(
                                   products: products,
                                   subspaces: subspaces,
                                   isCreate: false,
@@ -271,6 +269,7 @@ class AssignTagModelsDialog extends StatelessWidget {
                                   spaceTagModels: state.tags,
                                   onUpdate: (tags, subspaces) {
                                     onUpdate?.call(state.tags, subspaces);
+                                    Navigator.of(context).pop();
                                   },
                                 ),
                               );
@@ -318,10 +317,8 @@ class AssignTagModelsDialog extends StatelessWidget {
                                         (t) => t.internalId == tag.internalId);
                                   }
                                 }
-                                print("tryinh yo save");
-                              
+
                                 onUpdate?.call(state.tags, subspaces);
-                                
                               }
                             : null,
                         child: const Text('Save'),
diff --git a/lib/pages/spaces_management/space_model/widgets/tag_chips_display_widget.dart b/lib/pages/spaces_management/space_model/widgets/tag_chips_display_widget.dart
index 1f7e0d40..18efd501 100644
--- a/lib/pages/spaces_management/space_model/widgets/tag_chips_display_widget.dart
+++ b/lib/pages/spaces_management/space_model/widgets/tag_chips_display_widget.dart
@@ -94,21 +94,20 @@ class TagChipDisplay extends StatelessWidget {
                           barrierDismissible: false,
                           context: navigatorContext,
                           builder: (context) => AssignTagModelsDialog(
-                            products: products,
-                            subspaces: subspaces,
-                            allTags: allTags,
-                            initialTags: TagHelper.generateInitialTags(
-                                subspaces: subspaces,
-                                spaceTagModels: spaceModel?.tags ?? []),
-                            title: 'Edit Device',
-                            addedProducts:
-                                TagHelper.createInitialSelectedProducts(
-                                    spaceModel?.tags ?? [], subspaces),
-                            spaceName: spaceModel?.modelName ?? '',
-                            onUpdate: (tags, subspaces){
-                              print("here");
-                                onLoad?.call(tags, subspaces);}
-                          ),
+                              products: products,
+                              subspaces: subspaces,
+                              allTags: allTags,
+                              initialTags: TagHelper.generateInitialTags(
+                                  subspaces: subspaces,
+                                  spaceTagModels: spaceModel?.tags ?? []),
+                              title: 'Edit Device',
+                              addedProducts:
+                                  TagHelper.createInitialSelectedProducts(
+                                      spaceModel?.tags ?? [], subspaces),
+                              spaceName: spaceModel?.modelName ?? '',
+                              onUpdate: (tags, subspaces) {
+                                onLoad?.call(tags, subspaces);
+                              }),
                         );
                       }
                     },
@@ -139,7 +138,12 @@ class TagChipDisplay extends StatelessWidget {
                   allTags: allTags,
                   spaceName: spaceNameController.text,
                   isCreate: true,
-                  onLoad: (tags, subspaces) => onLoad?.call(tags, subspaces),
+                  onUpdate: (tags, subspaces) {
+                    onLoad?.call(tags, subspaces);
+                  },
+                  onLoad: (tags, subspaces) {
+                    onLoad?.call(tags, subspaces);
+                  },
                 ),
               );
             },
diff --git a/lib/pages/spaces_management/tag_model/views/add_device_type_model_widget.dart b/lib/pages/spaces_management/tag_model/views/add_device_type_model_widget.dart
index e324fa06..7944eb4f 100644
--- a/lib/pages/spaces_management/tag_model/views/add_device_type_model_widget.dart
+++ b/lib/pages/spaces_management/tag_model/views/add_device_type_model_widget.dart
@@ -41,12 +41,7 @@ class AddDeviceTypeModelWidget extends StatelessWidget {
 
   @override
   Widget build(BuildContext context) {
-    
-    if (spaceTagModels != null) {
-      for (var tag in spaceTagModels!) {
-        print(tag.tag);
-      }
-    }
+  
 
     final size = MediaQuery.of(context).size;
     final crossAxisCount = size.width > 1200
@@ -149,7 +144,7 @@ class AddDeviceTypeModelWidget extends StatelessWidget {
                                       initialTags: state.initialTag,
                                       title: dialogTitle,
                                       onUpdate: (tags, subspaces) {
-                                        onLoad?.call(tags, subspaces);
+                                        onUpdate?.call(tags, subspaces);
                                       },
                                     ),
                                   );