Few more fixes for the last commit