Fixing bug with missing units on upgrade.