Fix bugs in previous 2 commits