Fixed bug in main script