Fixed a few bugs in the  model