Fixes #19 Ensure balance is set after units