diff --git a/umap/static/umap/js/modules/rendering/ui.js b/umap/static/umap/js/modules/rendering/ui.js
index cd23ee31..4a3d584f 100644
--- a/umap/static/umap/js/modules/rendering/ui.js
+++ b/umap/static/umap/js/modules/rendering/ui.js
@@ -29,11 +29,9 @@ const FeatureMixin = {
 
   onRemove: function (map) {
     this.parentClass.prototype.onRemove.call(this, map)
-    if (map.editedFeature === this.feature) {
+    if (map._umap.editedFeature === this.feature) {
       this.feature.endEdit()
-      if (map.editedFeature === this.feature) {
-        map.editPanel.close()
-      }
+      map._umap.editPanel.close()
     }
   },