diff --git a/src/rbm/cd_train.py b/src/rbm/cd_train.py
index 669445b..31ddf22 100644
--- a/src/rbm/cd_train.py
+++ b/src/rbm/cd_train.py
@@ -10,21 +10,15 @@ class CdTrain:
 		self.state = state
 		self.params = params
 
-	def v_to_h(self, visible: np.ndarray) -> np.ndarray:
-		return np.vecmat(visible, self.state.w_hv) + self.state.b_h
-
-	def h_to_v(self, hidden: np.ndarray) -> np.ndarray:
-		return np.vecmat(hidden, np.transpose(self.state.w_hv)) + self.state.b_v
-
 	def v_to_h_prob(self, v: np.ndarray) -> np.ndarray:
-		state = self.v_to_h(v)
+		state = self.state.v_to_h(v)
 		if self.params.do_gaussian_visible:
 			return state
 
 		return prob(state)
 
 	def h_to_v_prob(self, h: np.ndarray) -> np.ndarray:
-		state = self.h_to_v(h)
+		state = self.state.h_to_v(h)
 		if self.params.do_gaussian_visible:
 			return state
 
diff --git a/src/rbm/layer.py b/src/rbm/layer.py
index e750911..18cd0e3 100644
--- a/src/rbm/layer.py
+++ b/src/rbm/layer.py
@@ -20,7 +20,7 @@ class RbmLayer:
 	def train_batch(self, v_states: np.ndarray, cd_func: Callable):
 		for epochs in range(self.params.num_epochs):
 			# Contrastive divergence learning: calculate gradients
-			cd_func(v_states, dwhv, dbh, dbv);
+			dwhv, dbh, dbv = cd_func(v_states)
 
 
 
@@ -36,9 +36,9 @@ if __name__ == "__main__":
 	v0 = uniform((1,2))
 	h0 = uniform((1,3))
 
-	h1 = l1.v_to_h(v0)
+	h1 = l1.state.v_to_h(v0)
 	print(h1.shape)
-	v1 = l1.h_to_v(h1)
+	v1 = l1.state.h_to_v(h1)
 	s_v1 = sample(v1)
 	p_h = prob(s_v1)
 
diff --git a/src/rbm/state.py b/src/rbm/state.py
index 3b71cd0..cc90f74 100644
--- a/src/rbm/state.py
+++ b/src/rbm/state.py
@@ -31,6 +31,12 @@ class RbmState:
 		obj = cls(w_hv, b_v, b_h)
 		return obj
 
+	def v_to_h(self, visible: np.ndarray) -> np.ndarray:
+		return np.vecmat(visible, self.w_hv) + self.b_h
+
+	def h_to_v(self, hidden: np.ndarray) -> np.ndarray:
+		return np.vecmat(hidden, np.transpose(self.w_hv)) + self.b_v
+
 	def to_file(self, filename: str):
 		np.savez(filename, whv=self.w_hv, bv=self.b_v, bh=self.b_h)
 		print(f"{filename} saved successfully!")