Small changes to the Main.java file