Small changes to main class