fixing bugs in flatten