diff --git a/src/java-jpa-entity.js b/src/java-jpa-entity.js
index 27606e7..b74c5a0 100644
--- a/src/java-jpa-entity.js
+++ b/src/java-jpa-entity.js
@@ -12,7 +12,7 @@ module.exports = {
             model.dir = model.package.replace(/\./g, '/')
             model.fields = model.fields.map(field => {
                 field.name = utils.camelcase(field.name)
-                field.type = javaType.mapJavaType(field.type, namespace)
+                field.type = javaType.mapJavaType(field.type, namespace, 'Entity')
 
                 if(field.type.fullQualifiedName && field.type.fullQualifiedName.includes('models')) {
                     field.type.fullQualifiedName = field.type.fullQualifiedName.replace('models', 'db')
diff --git a/src/java-type.js b/src/java-type.js
index f4009ae..87a25fd 100644
--- a/src/java-type.js
+++ b/src/java-type.js
@@ -50,19 +50,17 @@ const SUPPORTED_TYPES = {
 }
 
 
-function getCustomType(type, namespace) {
-
-console.log("====== CUSTOM: " + type)
+function getCustomType(type, namespace, suffix) {
 
     let match = LAST_ELEMENT_REGEX.exec(type)
     let name = ''
     let importstatement = null
     if(match) {
         name = utils.pascalcase(match[2])
-        importstatement = `${match[1].replace('enums', 'models')}.${name}`
+        importstatement = `${match[1].replace('enums', 'models')}.${name}${suffix}`
     }else {
         name = utils.pascalcase(type)
-        importstatement = `${namespace}.models.${name}`
+        importstatement = `${namespace}.models.${name}${suffix}`
     }
     
     return {
@@ -86,9 +84,9 @@ function mapCollection(type) {
 }
 
 module.exports = {
-    mapJavaType: (type, namespace) => {
+    mapJavaType: (type, namespace, suffix) => {
         const baseType = COLLECTION_REGEX.exec(type)[1]
-        let javaType = SUPPORTED_TYPES[baseType] || getCustomType(baseType, namespace)
+        let javaType = SUPPORTED_TYPES[baseType] || getCustomType(baseType, namespace, suffix)
 
         if(type.startsWith('[')) {
             javaType = mapCollection(javaType);