A few bug fixes to my last commit