Fix element removal bug