Fix to upgrade script