Fixes to bugs with previous commit