Fix bug by using model association instead of string comparison

Fix bug by using model association instead of string comparison #945