diff --git a/src/tests/test_label.py b/src/tests/test_label.py
index c546c42..7150872 100644
--- a/src/tests/test_label.py
+++ b/src/tests/test_label.py
@@ -10,7 +10,7 @@ if __name__ == "__main__":
 	prj_root = "/home/jens/work/repos/Rbm"
 
 	# Create Labeler
-	enc = Label(voc_size, 16, work_dir)
+	enc = Label(voc_size, 16, Label.EncodingType.OneHot, work_dir)
 
 	# Load
 	enc.fitter.load()
diff --git a/src/tests/test_learn_norbs_labels.py b/src/tests/test_learn_norbs_labels.py
index c8238f8..d27d06c 100644
--- a/src/tests/test_learn_norbs_labels.py
+++ b/src/tests/test_learn_norbs_labels.py
@@ -1,6 +1,5 @@
 import os
 import matplotlib.pyplot as plt
-
 from rbm.model import Model
 from rbm.entity import Entity, EntityParams, TrainingParams
 from rbm.matrix import Mat, np, read_armadillo
@@ -10,7 +9,7 @@ from rbm.label import Label
 class TestModel(Model):
 	def __init__(self, name: str, work_dir: str = '.'):
 		super().__init__(name, work_dir)
-#		self.unit_image = Entity((96*96, 16), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False), TrainingParams(learning_rate=0.00001, momentum=0.9, do_rao_blackwell=True, num_epochs=1000))
+#		self.unit_image = Entity((96*96, 16), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False, num_gibbs_samples=3), TrainingParams(learning_rate=0.00001, momentum=0.9, do_rao_blackwell=False, num_epochs=1000))
 		self.unit_image = Entity((96 * 96, 16), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False))
 		self.unit_combine = Entity((32, 64), EntityParams(), TrainingParams(learning_rate=0.01, momentum=0.9, do_rao_blackwell=True, num_epochs=1000))
 
@@ -55,7 +54,7 @@ if __name__ == "__main__":
 
 	# Train
 	# Create Labeler
-	enc = Label(20, 16, work_dir)
+	enc = Label(20, 16, Label.EncodingType.OneHot, work_dir)
 
 	# Load
 	enc.fitter.load()
@@ -71,19 +70,20 @@ if __name__ == "__main__":
 	# save state
 	model.save()
 
-	fig, axes = plt.subplots(3, len(test_batch), figsize=(12, 3))
-	for index, image in enumerate(test_batch):
+	cmap = 'Grays'
+	fig, axes = plt.subplots(3, len(train_batch), figsize=(12, 3))
+	for index, image in enumerate(train_batch):
 		label_zero = np.zeros(shape=16)
 		image_zero = np.zeros(shape=96*96)
 		image_reconst, labels_reconst = model.reconstruct(model.forward2(image, label_zero))
 #		image_reconst, labels_reconst = model.reconstruct(model.forward2(image_zero, encoded[index, :]))
-		img = 2*(image_reconst + 0.5)
+		img = 1-2*(image_reconst + 0.5)
 		img = np.reshape(img, (96, 96))
-		axes[0][index].imshow(img)
+		axes[0][index].imshow(img, cmap=cmap)
 		axes[0][index].axis('off')
-		axes[1][index].imshow(np.reshape(labels_reconst, (4,4)))
+		axes[1][index].imshow(np.reshape(labels_reconst, (4,4)), cmap=cmap)
 		axes[1][index].axis('off')
-		axes[2][index].imshow(np.reshape(encoded[index, :], (4,4)))
+		axes[2][index].imshow(np.reshape(encoded[index, :], (4,4)), cmap=cmap)
 		axes[2][index].axis('off')
 
 	plt.show()