diff --git a/source/Layer.hpp b/source/Layer.hpp
index 530137b..a3760c5 100644
--- a/source/Layer.hpp
+++ b/source/Layer.hpp
@@ -32,26 +32,17 @@ public:
 
 	Json::Value toJson() const;
 
-	void setBatch(arma::mat const &batch)
+	void calcContextBatch(arma::mat &batch)
 	{
-		if (batch.n_rows > 0)
-		{	
-			if (m_numContext)
+		size_t numTraining = batch.n_rows;
+
+		if (m_numContext > 0 and numTraining > 0)
+		{
+			for (int i=1; i < numTraining; i++)
 			{
-				m_context = arma::zeros(batch.n_rows, m_numContext);
-				arma::mat ctx = arma::zeros(1, m_numContext);
-				for (int i=1; i < batch.n_rows; i++)
-				{
-					arma::mat h = prob(v_to_h(arma::join_rows(batch.row(i-1), ctx)));
-					ctx = h;
-					m_context.row(i) = ctx;
-				}
-				arma::mat batch_with_ctx = arma::join_rows(batch, m_context);
-				Rbm::setBatch(trainingData(batch_with_ctx));
-			}
-			else
-			{
-				Rbm::setBatch(trainingData(batch));
+				arma::mat h = toHiddenProbs(batch.row(i-1));
+				arma::mat training_with_ctx = arma::join_rows(batch.row(i-1).cols(0, numVisible()-m_numContext-1), h);
+				batch.row(i-1) = training_with_ctx;
 			}
 		}
 	}