diff --git a/source/MainComponent.cpp b/source/MainComponent.cpp
index 066651d..292b449 100644
--- a/source/MainComponent.cpp
+++ b/source/MainComponent.cpp
@@ -505,7 +505,7 @@ void MainComponent::buttonClicked (Button* buttonThatWasClicked)
         //[UserButtonCode_ShakeButton] -- add your button handler code here..
     	m_pLayer->weightsInit(weightInitLabel->getText().getFloatValue());
 		m_pLayer->redrawWeights();
-    	m_pLayer->redrawReconstruction(m_stack->trainingData(m_pLayer).row(m_trainingIndex));
+    	m_pLayer->redrawReconstruction();
         //[/UserButtonCode_ShakeButton]
     }
     else if (buttonThatWasClicked == createButton)
@@ -641,7 +641,7 @@ void MainComponent::sliderValueChanged (Slider* sliderThatWasMoved)
     {
         //[UserSliderCode_numGibbsSlider] -- add your slider handling code here..
     	m_pLayer->params().numGibbs = sliderThatWasMoved->getValue();
-		m_pLayer->redrawReconstruction(m_stack->trainingData().row(m_trainingIndex));
+		m_pLayer->redrawReconstruction();
         //[/UserSliderCode_numGibbsSlider]
     }
     else if (sliderThatWasMoved == m_progressBarSlider)
@@ -841,7 +841,7 @@ bool MainComponent::onProgress(Rbm *pRbm, const Rbm::Status &status)
 {
 	RbmComponent *pComp = static_cast<RbmComponent*>(pRbm);
 	pComp->upPass(pComp->DrawTraining->getData());
-	pComp->redrawReconstruction(pComp->DrawTraining->getData());
+	pComp->redrawReconstruction();
 	pComp->redrawWeights();
 
 	m_progressBarSlider->setValue(status.progress);
diff --git a/source/RbmComponent.cpp b/source/RbmComponent.cpp
index c1732c7..7d7042c 100644
--- a/source/RbmComponent.cpp
+++ b/source/RbmComponent.cpp
@@ -218,20 +218,15 @@ void RbmComponent::buttonClicked(Button* buttonThatWasClicked)
 			RbmComponent *pComp = static_cast<RbmComponent*> (next);
 			pComp->enable(state);
 		}
+		redrawReconstruction();
 	}
 }
 
-void RbmComponent::redrawReconstruction(const arma::mat& v)
+void RbmComponent::redrawReconstruction()
 {
-	uint32_t i;
+	RbmComponent *pComp = static_cast<RbmComponent*> (root());
+	pComp->upDownPass(pComp->DrawTraining->getData());
 
-	for (i=0; i < params().numGibbs; i++)
-	{
-		DrawHidden->getData() = toHiddenProbs(v);
-		DrawHidden->DrawData();
-		DrawReconstruction->getData() = toVisibleProbs(DrawHidden->getData());
-		DrawReconstruction->DrawData();
-	}
 }
 
 void RbmComponent::upPass(const arma::mat& v)
@@ -322,7 +317,8 @@ void RbmComponent::redrawWeights(size_t index)
 
 void RbmComponent::setTrainingData(arma::mat const& batch)
 {
-	upDownPass(trainingData(batch));
+	RbmComponent *pComp = static_cast<RbmComponent*> (root());
+	pComp->upDownPass(batch);
 }
 //[/MiscUserCode]
 
diff --git a/source/RbmComponent.hpp b/source/RbmComponent.hpp
index 87380e7..40603df 100644
--- a/source/RbmComponent.hpp
+++ b/source/RbmComponent.hpp
@@ -64,7 +64,7 @@ public:
 	void redrawWeights();
 	void redrawWeights(size_t index);
 
-	void redrawReconstruction(arma::mat const &v);
+	void redrawReconstruction();
 	
 	void upPass(arma::mat const &v);
 	void downPass(arma::mat const &v);