minor bug fix

Minor bug fix in the models