diff --git a/source/MainComponent.cpp b/source/MainComponent.cpp
index 7f2a939..32a1e7c 100644
--- a/source/MainComponent.cpp
+++ b/source/MainComponent.cpp
@@ -498,7 +498,8 @@ void MainComponent::buttonClicked (Button* buttonThatWasClicked)
     else if (buttonThatWasClicked == addButton)
     {
         //[UserButtonCode_addButton] -- add your button handler code here..
-//		m_layers.add(m_pRbmComponent[0]->getTrainingData(), m_weights[0]->getNumVisible());
+		RbmComponent *pComp = static_cast<RbmComponent*>(m_stack->getLayer(0));
+		addTraining(pComp->DrawTraining->getData());
         //[/UserButtonCode_addButton]
     }
     else if (buttonThatWasClicked == ShakeButton)
@@ -512,7 +513,16 @@ void MainComponent::buttonClicked (Button* buttonThatWasClicked)
     else if (buttonThatWasClicked == createButton)
     {
         //[UserButtonCode_createButton] -- add your button handler code here..
-		create("justCreated");
+		int numVisX = numVisibleLabel->getText().getIntValue();
+		int numVisY = numVisibleYLabel->getText().getIntValue();
+		int numHid = numHiddenLabel->getText().getIntValue();
+		m_stack = new Stack(std::string(projectNameLabel->getText().getCharPointer()));
+		m_stack->addLayer(onConstruct("Layer", 0, numVisX, numVisY, numHid));
+		m_rbmSelect->clear(dontSendNotification);
+		m_rbmSelect->addItem(String(0), 0+1);
+		m_rbmSelect->setSelectedId(1, sendNotification);
+//		addTraining(arma::zeros(1, numVisX*numVisY));
+		
         //[/UserButtonCode_createButton]
     }
     else if (buttonThatWasClicked == loadButton)
diff --git a/source/MainComponent.hpp b/source/MainComponent.hpp
index 50024a2..4f742ab 100644
--- a/source/MainComponent.hpp
+++ b/source/MainComponent.hpp
@@ -77,7 +77,7 @@ public:
 	{
 		RbmComponent *pComp = new RbmComponent(name, id, numVisibleX, numVisibleY, numHidden, this);
 		addAndMakeVisible(pComp);
-		return pComp;
+		return static_cast<Layer*>(pComp);
 	}
 
 	
diff --git a/source/RbmComponent.hpp b/source/RbmComponent.hpp
index 8574a2b..ef16721 100644
--- a/source/RbmComponent.hpp
+++ b/source/RbmComponent.hpp
@@ -134,11 +134,11 @@ public:
 	arma::mat const& getTopWeights() override;
 	arma::mat getConvolutedWeight(arma::mat const &h) override;
 	arma::mat const& getWeights() override;
+    ScopedPointer<DrawComponent> DrawTraining;
 	
 private:
     //[UserVariables]   -- You can add your own custom variables in this section.
 	RbmComponentListener *m_listener;
-    ScopedPointer<DrawComponent> DrawTraining;
     ScopedPointer<DrawComponent> DrawReconstruction;
     ScopedPointer<DrawComponent> DrawWeights;
     ScopedPointer<DrawComponent> DrawHidden;