diff --git a/src/rbm/cd_train.py b/src/rbm/cd_train.py
index eb27117..0cbf1bf 100644
--- a/src/rbm/cd_train.py
+++ b/src/rbm/cd_train.py
@@ -2,10 +2,10 @@ import numpy as np
 from collections.abc import Callable
 
 from params import RbmParams
-from matrix import sample, gaussian
+from matrix import prob, sample, gaussian
 
 def cd_jens(v_states: np.ndarray, params: RbmParams, v_to_ph: Callable, h_to_pv: Callable):
-	v_probs = v_states
+	v_probs = prob(v_states)
 	h_states = v_to_ph(v_states)
 	h_probs = h_states
 
diff --git a/src/rbm/entity.py b/src/rbm/entity.py
index 9ff2cea..50bc93a 100644
--- a/src/rbm/entity.py
+++ b/src/rbm/entity.py
@@ -76,18 +76,18 @@ class Entity:
 
 
 	def gibbs_v_to_h(self, v: np.ndarray) -> np.ndarray:
-		h = None
-		for i in range(self.params.num_gibbs_samples):
-			h = self.v_to_ph(v)
-			v = self.h_to_pv(h)
+		h = self.v_to_ph(v)
+		for i in range(self.params.num_gibbs_samples-1):
+			h = self.h_to_pv(h)
+			h = self.v_to_ph(h)
 
 		return h
 
 	def gibbs_h_to_v(self, h: np.ndarray) -> np.ndarray:
-		v = None
-		for i in range(self.params.num_gibbs_samples):
-			v = self.h_to_pv(h)
-			h = self.v_to_ph(v)
+		v = self.h_to_pv(h)
+		for i in range(self.params.num_gibbs_samples-1):
+			v = self.v_to_ph(v)
+			v = self.h_to_pv(v)
 
 		return v