Minor modification to models