diff --git a/src/rbm/params.py b/src/rbm/params.py
index 42eb796..7525947 100644
--- a/src/rbm/params.py
+++ b/src/rbm/params.py
@@ -11,21 +11,23 @@ class Params:
 
 class EntityParams(Params):
 	def __init__(self):
+		# Training parameters
 		self.learning_rate = 0.1
 		self.momentum = 0.5
 		self.weight_decay = 0
 		self.num_epochs = 1000
-		self.num_gibbs_samples = 1
 		self.mini_batch_size = 0
-
-		self.do_gaussian_visible = False
-		self.do_gaussian_hidden = False
 		self.do_rao_blackwell = False
-
 		self.do_gibbs_sample_visible = False
 		self.do_gibbs_sample_hidden = False
 		self.do_batch_sample = False
 
+		# Entity parameters
+		self.num_gibbs_samples = 1
+		self.do_gaussian_visible = False
+		self.do_gaussian_hidden = False
+
+
 	@classmethod
 	def from_dict(cls, params: dict, version: str = '0'):
 		obj = EntityParams()