BUG: updated based on bug fix in previous commit