Fix forgetting to put back code I took away