add return to main