Modified model.py file