fixed minor error in main
Merged the changes on main.cpp