Minor code adjustments for better output