Fixed some bugs

Found some bugs while doing a quick code review