Fixed the entity issue