FIX forgot to add models.py changes