diff --git a/generator/src/main/java/org/onosproject/yangutils/translator/tojava/TempJavaFragmentFiles.java b/generator/src/main/java/org/onosproject/yangutils/translator/tojava/TempJavaFragmentFiles.java
index 7f6eb84..abd72eb 100644
--- a/generator/src/main/java/org/onosproject/yangutils/translator/tojava/TempJavaFragmentFiles.java
+++ b/generator/src/main/java/org/onosproject/yangutils/translator/tojava/TempJavaFragmentFiles.java
@@ -116,8 +116,10 @@
 import static org.onosproject.yangutils.utils.UtilConstants.INTERFACE;
 import static org.onosproject.yangutils.utils.UtilConstants.INVOCATION_TARGET_EXCEPTION;
 import static org.onosproject.yangutils.utils.UtilConstants.INVOCATION_TARGET_EXCEPTION_IMPORT;
+import static org.onosproject.yangutils.utils.UtilConstants.ITR_IMPORT;
 import static org.onosproject.yangutils.utils.UtilConstants.JAVA_UTIL_PKG;
 import static org.onosproject.yangutils.utils.UtilConstants.METHOD;
+import static org.onosproject.yangutils.utils.UtilConstants.METHOD_IMPORT;
 import static org.onosproject.yangutils.utils.UtilConstants.NEW_LINE;
 import static org.onosproject.yangutils.utils.UtilConstants.OPERATION_TYPE_ATTRIBUTE;
 import static org.onosproject.yangutils.utils.UtilConstants.OPERATION_TYPE_CLASS;
@@ -128,6 +130,7 @@
 import static org.onosproject.yangutils.utils.UtilConstants.REFLECT_IMPORTS;
 import static org.onosproject.yangutils.utils.UtilConstants.SELECT_LEAF;
 import static org.onosproject.yangutils.utils.UtilConstants.SERVICE;
+import static org.onosproject.yangutils.utils.UtilConstants.SET_IMPORT;
 import static org.onosproject.yangutils.utils.UtilConstants.SLASH;
 import static org.onosproject.yangutils.utils.UtilConstants.SUBTREE_FILTERED;
 import static org.onosproject.yangutils.utils.UtilConstants.VALUE_LEAF;
@@ -1788,6 +1791,9 @@
         }
 
         if (curNode.isOpTypeReq()) {
+            if (curNode instanceof YangAugmentableNode) {
+                addImportsForSubTreeFilterAug(imports);
+            }
             addSubTreeImportStrings(imports);
         }
 
@@ -1836,6 +1842,18 @@
         freeTemporaryResources(false);
     }
 
+    private void addImportsForSubTreeFilterAug(List<String> imports) {
+        if (!imports.contains(SET_IMPORT)) {
+            imports.add(SET_IMPORT);
+        }
+        if (!imports.contains(METHOD_IMPORT)) {
+            imports.add(METHOD_IMPORT);
+        }
+        if (!imports.contains(ITR_IMPORT)) {
+            imports.add(ITR_IMPORT);
+        }
+    }
+
     //Adds import for array list.
     private void addArrayListImport(List<String> imports) {
         if (imports.contains(javaImportData.getImportForList())) {
diff --git a/generator/src/main/java/org/onosproject/yangutils/translator/tojava/utils/SubtreeFilteringMethodsGenerator.java b/generator/src/main/java/org/onosproject/yangutils/translator/tojava/utils/SubtreeFilteringMethodsGenerator.java
index 2a1b28f..a02155b 100644
--- a/generator/src/main/java/org/onosproject/yangutils/translator/tojava/utils/SubtreeFilteringMethodsGenerator.java
+++ b/generator/src/main/java/org/onosproject/yangutils/translator/tojava/utils/SubtreeFilteringMethodsGenerator.java
@@ -932,35 +932,104 @@
         return builder.toString();
     }
 
-    //Returns method string for op params augmented syntax
-    static String getAugmentableSubTreeFiltering() {
-        return "        for (Object augmentInfo : yangAugmentedInfoMap()" +
-                ".values()) {\n" +
-                "            Object appInstanceInfo = appInstance.yangAugmentedInfo(" +
-                "augmentInfo.getClass());\n" +
-                "            if (appInstanceInfo == null) {\n" +
-                "                subTreeFilteringResultBuilder.addYangAugmentedInfo(" +
-                "augmentInfo, augmentInfo.getClass());\n" +
-                "            } else {\n" +
-                "                Object processSubtreeFiltering;\n" +
-                "                try {\n" +
-                "                    Class<?> augmentedClass = augmentInfo" +
-                ".getClass();\n" +
-                "                    processSubtreeFiltering = augmentInfo.getClass()" +
-                ".getMethod(\"processSubtreeFiltering\", augmentedClass).invoke(" +
-                "augmentInfo, appInstanceInfo);\n" +
-                "                    if (processSubtreeFiltering != null) {\n" +
-                "                        subTreeFilteringResultBuilder" +
-                ".addYangAugmentedInfo(processSubtreeFiltering, processSubtreeFiltering.getClass());\n" +
+    public static String getAugmentableSubTreeFiltering() {
+        return "        if (yangAugmentedInfoMap.isEmpty()) {\n            " +
+                "Set<Map.Entry<Class<?>, Object>> augment = appInstance" +
+                ".yangAugmentedInfoMap().entrySet();\n            " +
+                "if (augment != null && !augment.isEmpty()) {\n" +
+                "                " +
+                "Iterator<Map.Entry<Class<?>, Object>> augItr = " +
+                "augment.iterator();\n                " +
+                "while (augItr.hasNext()) {\n                    " +
+                "Map.Entry<Class<?>, Object> aug = augItr.next();\n" +
+                "                    " +
+                "Class<?> augClass = aug.getKey();\n                    " +
+                "String augClassName = augClass.getName();\n" +
+                "                    " +
+                "int index = augClassName.lastIndexOf('.');\n" +
+                "                    " +
+                "String classPackage = augClassName.substring(0, index) +\n" +
+                "                            " +
+                "\".\" + \"Default\" + augClass.getSimpleName() + \"$\"\n" +
+                "                            " +
+                "+ augClass.getSimpleName() + \"Builder\";\n" +
+                "                    " +
+                "ClassLoader classLoader = augClass.getClassLoader();\n" +
+                "                    " +
+                "try {\n                        " +
+                "Class<?> builderClass;\n                        " +
+                "builderClass = classLoader.loadClass(classPackage);\n" +
+                "                        " +
+                "Object builderObj = builderClass.newInstance();\n" +
+                "                        " +
+                "Method method = builderClass.getMethod(\"build\");\n" +
+                "                        " +
+                "Object defaultObj = method.invoke(builderObj);\n" +
+                "                        " +
+                "Class<?> defaultClass = defaultObj.getClass();\n" +
+                "                        " +
+                "method = defaultClass.getMethod\n" +
+                "                                " +
+                "(\"processSubtreeFiltering\", augClass,\n" +
+                "                                 " +
+                "boolean.class);\n                        " +
+                "Object result = method.invoke(defaultObj, aug.getValue(),\n" +
+                "                                                      " +
+                "true);\n                        " +
+                "subTreeFilteringResultBuilder\n" +
+                "                                " +
+                ".addYangAugmentedInfo(result, augClass);\n" +
+                "                    " +
+                "} catch (ClassNotFoundException | InstantiationException\n" +
+                "                            | NoSuchMethodException |\n" +
+                "                            " +
+                "InvocationTargetException | IllegalAccessException e) {\n" +
+                "                        e.printStackTrace();\n" +
                 "                    }\n" +
-                "                } catch (NoSuchMethodException |" +
-                " InvocationTargetException | IllegalAccessException e) {\n" +
-                "                    continue;\n" +
+                "                }\n" +
+                "            }\n" +
+                "        } else {\n            " +
+                "Set<Map.Entry<Class<?>, Object>> augment = " +
+                "yangAugmentedInfoMap\n                    .entrySet();\n" +
+                "            " +
+                "Iterator<Map.Entry<Class<?>, Object>> augItr = " +
+                "augment.iterator();\n            " +
+                "while (augItr.hasNext()) {\n                " +
+                "Map.Entry<Class<?>, Object> aug = augItr.next();\n" +
+                "                Class<?> augClass = aug.getKey();\n" +
+                "                " +
+                "Object appInstanceInfo = appInstance.yangAugmentedInfo(" +
+                "augClass);\n                if (appInstanceInfo == null) {\n" +
+                "                    " +
+                "subTreeFilteringResultBuilder.addYangAugmentedInfo\n" +
+                "                            " +
+                "(aug.getValue(), aug.getKey());\n" +
+                "                } else {\n                    " +
+                "Object processSubtreeFiltering;\n                    try {\n" +
+                "                        " +
+                "processSubtreeFiltering = aug.getValue().getClass()\n" +
+                "                                " +
+                ".getMethod(\"processSubtreeFiltering\",\n" +
+                "                                           " +
+                "aug.getKey(), boolean.class)\n" +
+                "                                .invoke(aug.getValue(),\n" +
+                "                                        " +
+                "appInstanceInfo, true);\n                        " +
+                "if (processSubtreeFiltering != null) {\n" +
+                "                            " +
+                "subTreeFilteringResultBuilder\n                            " +
+                "        .addYangAugmentedInfo(processSubtreeFiltering, " +
+                "aug.getKey());\n                        }\n" +
+                "                    } catch (NoSuchMethodException | " +
+                "InvocationTargetException | IllegalAccessException e) {\n" +
+                "                        e.printStackTrace();\n" +
+                "                    }\n" +
                 "                }\n" +
                 "            }\n" +
                 "        }\n";
     }
 
+
     private static String getSubTreeFilteredCondition(String name) {
         StringBuilder builder = new StringBuilder();
         String cond = SUBTREE_FILTERED + SPACE + AND_OPERATION + SPACE + name +
diff --git a/generator/src/main/java/org/onosproject/yangutils/utils/UtilConstants.java b/generator/src/main/java/org/onosproject/yangutils/utils/UtilConstants.java
index 3a1a1f5..025c8fc 100644
--- a/generator/src/main/java/org/onosproject/yangutils/utils/UtilConstants.java
+++ b/generator/src/main/java/org/onosproject/yangutils/utils/UtilConstants.java
@@ -1691,6 +1691,23 @@
             IMPORT + COLLECTION_IMPORTS + ".ArrayList;\n";
 
     /**
+     * Static attribute for set collection import.
+     */
+    public static final String SET_IMPORT = IMPORT + COLLECTION_IMPORTS +
+            PERIOD + SET + SEMI_COLON + NEW_LINE;
+
+    /**
+     * Static attribute for reflect method import.
+     */
+    public static final String METHOD_IMPORT = IMPORT + REFLECT_IMPORTS +
+            PERIOD + METHOD + SEMI_COLON + NEW_LINE;
+
+    /**
+     * Static attribute for collection iterator import.
+     */
+    public static final String ITR_IMPORT = IMPORT + COLLECTION_IMPORTS +
+            PERIOD + "Iterator;\n";
+    /**
      * Static attribute for unused keyword.
      */
     public static final String UNUSED = "UNUSED";