Edit mode fix