diff --git a/src/java-jpa-entity.js b/src/java-jpa-entity.js
index 165544f..27606e7 100644
--- a/src/java-jpa-entity.js
+++ b/src/java-jpa-entity.js
@@ -14,8 +14,8 @@ module.exports = {
                 field.name = utils.camelcase(field.name)
                 field.type = javaType.mapJavaType(field.type, namespace)
 
-                if(field.type.fullQualifiedName && field.type.fullQualifiedName.includes('model') && !field.type.isEnum) {
-                    field.type.fullQualifiedName = field.type.fullQualifiedName.replace('model', 'db')
+                if(field.type.fullQualifiedName && field.type.fullQualifiedName.includes('models')) {
+                    field.type.fullQualifiedName = field.type.fullQualifiedName.replace('models', 'db')
                 }
 
                 field.default = javaType.mapDefaultValue(field.type, field.default)