Flatten by removing unnecessary class