Fixes bugs with upgrade process