Another bug fix to script