Minor bug fix

moved function to correct model