Misc code fix