Minor fixes to balance checking