diff --git a/utils/yangutils/src/main/java/org/onosproject/yangutils/translator/tojava/javamodel/YangJavaModule.java b/utils/yangutils/src/main/java/org/onosproject/yangutils/translator/tojava/javamodel/YangJavaModule.java
index 235ebc5..d6ba066 100644
--- a/utils/yangutils/src/main/java/org/onosproject/yangutils/translator/tojava/javamodel/YangJavaModule.java
+++ b/utils/yangutils/src/main/java/org/onosproject/yangutils/translator/tojava/javamodel/YangJavaModule.java
@@ -23,6 +23,7 @@
 import org.onosproject.yangutils.translator.tojava.JavaImportData;
 import org.onosproject.yangutils.translator.tojava.TempJavaCodeFragmentFiles;
 import org.onosproject.yangutils.translator.tojava.utils.YangJavaModelUtils;
+import org.onosproject.yangutils.translator.tojava.utils.YangPluginConfig;
 
 import static org.onosproject.yangutils.translator.tojava.GeneratedJavaFileType.GENERATE_INTERFACE_WITH_BUILDER;
 import static org.onosproject.yangutils.translator.tojava.utils.JavaIdentifierSyntax.getRootPackage;
@@ -126,13 +127,13 @@
     /**
      * Generates java code for module.
      *
-     * @param baseCodeGenDir code generation directory
+     * @param yangPlugin YANG plugin config
      * @throws IOException when fails to generate the source files
      */
     @Override
-    public void generateCodeEntry(String baseCodeGenDir) throws IOException {
+    public void generateCodeEntry(YangPluginConfig yangPlugin) throws IOException {
         String modulePkg = getRootPackage(getVersion(), getNameSpace().getUri(), getRevision().getRevDate());
-        YangJavaModelUtils.generateCodeOfRootNode(this, baseCodeGenDir, modulePkg);
+        YangJavaModelUtils.generateCodeOfRootNode(this, yangPlugin, modulePkg);
     }
 
     @Override