diff --git a/src/rbm/train.py b/src/rbm/train.py
index b166b7c..38fca06 100644
--- a/src/rbm/train.py
+++ b/src/rbm/train.py
@@ -214,12 +214,13 @@ def train(entity: Entity, batch: Mat, status: Status):
 	if entity.type == Entity.Type.BG_RBM:
 		cd_func = cd_binary_gaussian
 
-	for mini_batch in to_mini_batch(batch, mini_batch_size):
+	for epochs in range(params.num_epochs):
 		if not keep_running:
 			break
 
-		entity.grad_zero()
-		for epochs in range(params.num_epochs):
+		for mini_batch in to_mini_batch(batch, mini_batch_size):
+			# Reset gradient
+			entity.grad_zero()
 			# Contrastive divergence learning: calculate gradients
 			dwhv, dbv, dbh = cd_func(entity, mini_batch)