Small changes to main loop