Fixed some bugs introduced in previous commit