Minor changes in main()