Fixing minor issue with the script