Updated controller to make heroku work