Small bug fix (added some return statements)