Correct model associations