diff --git a/AlgebraicDataflowArchitectureModel/src/application/views/FlowLayerWindow.java b/AlgebraicDataflowArchitectureModel/src/application/views/FlowLayerWindow.java
index 4256204..d65d76e 100644
--- a/AlgebraicDataflowArchitectureModel/src/application/views/FlowLayerWindow.java
+++ b/AlgebraicDataflowArchitectureModel/src/application/views/FlowLayerWindow.java
@@ -4,12 +4,16 @@
 import application.editor.Editor;
 import application.editor.IStageChangeListener;
 import application.editor.Stage;
+import application.editor.stages.ControlFlowModelingStage;
 
 import javax.swing.*;
 import java.awt.*;
 import java.awt.event.ActionEvent;
 import java.awt.event.ActionListener;
 
+/**
+ * {@link FlowLayerWindow} shows the window for enabling / disabling layers in the graph.
+ */
 public class FlowLayerWindow extends JDialog implements IStageChangeListener {
 	private final String title = "Flow Layer";
 	private final JCheckBox dataFlowCheckBox;
@@ -71,15 +75,25 @@
 	}
 	
 	/**
-	 * Toggle the activation of layer check boxes whenever the current stage changes.
+	 * Enable all checkboxes when the stage is changed to {@link ControlFlowModelingStage} <br>
+	 * Disable all checkboxes whenever the active stage changed otherwise
 	 */
 	@Override
 	public void stageChanged(Stage newStage) {
-		dataFlowCheckBox.setEnabled(false);
-		pushFlowCheckBox.setEnabled(false);
-		pullFlowCheckBox.setEnabled(false);
-		
-		newStage.setLayerEnabled(Stage.PUSH_FLOW_LAYER, false);
-		newStage.setLayerEnabled(Stage.PULL_FLOW_LAYER, false);
+		if (newStage instanceof ControlFlowModelingStage) {
+			dataFlowCheckBox.setEnabled(true);
+			pushFlowCheckBox.setEnabled(true);
+			pullFlowCheckBox.setEnabled(true);
+			
+			newStage.setLayerEnabled(Stage.PUSH_FLOW_LAYER, pushFlowCheckBox.isSelected());
+			newStage.setLayerEnabled(Stage.PULL_FLOW_LAYER, pullFlowCheckBox.isSelected());
+		} else {
+			dataFlowCheckBox.setEnabled(false);
+			pushFlowCheckBox.setEnabled(false);
+			pullFlowCheckBox.setEnabled(false);
+			
+			newStage.setLayerEnabled(Stage.PUSH_FLOW_LAYER, false);
+			newStage.setLayerEnabled(Stage.PULL_FLOW_LAYER, false);
+		}
 	}
 }