diff --git a/src/java-jpa-entity.js b/src/java-jpa-entity.js
index 1ae5c8c..d93d8a6 100644
--- a/src/java-jpa-entity.js
+++ b/src/java-jpa-entity.js
@@ -16,8 +16,8 @@ module.exports = {
                 field.type = javaType.mapJavaType(field.type, namespace, field.isEnum ? '' : 'Record')
                 
                 if(field.type.fullQualifiedName && field.type.fullQualifiedName.includes('models') && !field.isEnum) {
-                    field.type.fullQualifiedName = field.type.fullQualifiedName.replace('models', 'db')
-
+                    field.type.fullQualifiedName = field.type.fullQualifiedName.replace('models', suffix)
+                    field.type.imports = field.type.imports.map(i => i.replace(`${namespace}.models`, `${namespace}.${suffix}`))
                 }
 
                 field.default = javaType.mapDefaultValue(field.type, field.default)