Fixing remove bug