Small model bug fix