Correct output extension when flatten is false