diff --git a/AlgebraicDataflowArchitectureModel/src/application/editor/Editor.java b/AlgebraicDataflowArchitectureModel/src/application/editor/Editor.java
index dd86d2a..7eadb4b 100644
--- a/AlgebraicDataflowArchitectureModel/src/application/editor/Editor.java
+++ b/AlgebraicDataflowArchitectureModel/src/application/editor/Editor.java
@@ -206,6 +206,10 @@
 
                 // Force to change PushPullSelectionStage to construct mxGraph
                 changeStage(STAGE_PUSH_PULL_SELECTION);
+                if (curStage instanceof PushPullSelectionStage) {
+                    PushPullSelectionStage stage = (PushPullSelectionStage) curStage;
+                    stage.constructGraph();
+                }
 
                 // Update current file info
                 curFilePath = file.getAbsolutePath();
@@ -240,6 +244,10 @@
 
                 // Force to change PushPullSelectionStage to construct mxGraph
                 changeStage(STAGE_PUSH_PULL_SELECTION);
+                if (curStage instanceof PushPullSelectionStage) {
+                    PushPullSelectionStage stage = (PushPullSelectionStage) curStage;
+                    stage.constructGraph();
+                }
 
                 // Update current file info
                 curFilePath = file.getAbsolutePath();
diff --git a/AlgebraicDataflowArchitectureModel/src/application/editor/stages/PushPullSelectionStage.java b/AlgebraicDataflowArchitectureModel/src/application/editor/stages/PushPullSelectionStage.java
index 64edd98..0dfe517 100644
--- a/AlgebraicDataflowArchitectureModel/src/application/editor/stages/PushPullSelectionStage.java
+++ b/AlgebraicDataflowArchitectureModel/src/application/editor/stages/PushPullSelectionStage.java
@@ -35,7 +35,6 @@
         if (prevStage instanceof DataFlowModelingStage) {
             model = prevStage.getModel();
             dataFlowGraph = analyzeDataTransferModel(model);
-            constructGraph();
         }
     }