Fixed some small bugs with last commit