diff --git a/source/Rbm.cpp b/source/Rbm.cpp
index fbe7c55..1c11ca3 100644
--- a/source/Rbm.cpp
+++ b/source/Rbm.cpp
@@ -73,13 +73,17 @@ Json::Value Rbm::toJson() const
 
 void Rbm::weightUpdate(arma::mat const &v_states, arma::mat &dw, arma::mat &dbh, arma::mat &dbv)
 {
-	arma::mat h_probs = prob(v_to_h(v_states));
 	arma::mat v_probs(dbv.n_rows, dbv.n_cols);
 
-	arma::mat h_states = h_probs;
+	arma::mat h_states = v_to_h(v_states);
+	arma::mat h_probs = prob(h_states);
 	
 	// Sample hidden
-	if (!m_params.doRaoBlackwell)
+	if (m_params.doRaoBlackwell)
+	{
+		h_states = h_probs;
+	}
+	else
 	{
 		h_states = sample(h_probs);
 	}