diff --git a/src/ModelGeneration.ts b/src/ModelGeneration.ts
index 143fd30..fe27e46 100644
--- a/src/ModelGeneration.ts
+++ b/src/ModelGeneration.ts
@@ -73,7 +73,7 @@ function removeUnusedImports(rendered: string) {
         .split(",");
     const restOfEntityDefinition = rendered.substring(closeBracketIndex);
     const distinctImports = imports.filter(
-        v => restOfEntityDefinition.indexOf(v) !== -1
+        v => restOfEntityDefinition.indexOf(`@${v}(`) !== -1
     );
     return `${rendered.substring(0, openBracketIndex)}${distinctImports.join(
         ","