diff --git a/src/tests/test_learn_norbs_labels.py b/src/tests/test_learn_norbs_labels.py
index 369431b..c8238f8 100644
--- a/src/tests/test_learn_norbs_labels.py
+++ b/src/tests/test_learn_norbs_labels.py
@@ -4,35 +4,34 @@ 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
+from rbm.train import train, Status
+from rbm.label import Label
 
 class TestModel(Model):
 	def __init__(self, name: str, work_dir: str = '.'):
 		super().__init__(name, work_dir)
-		self.unit1 = Entity((96*96, 16), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=True), TrainingParams(learning_rate=0.00001, momentum=0.9, do_rao_blackwell=True, num_epochs=1000, num_gibbs_samples=1))
-		self.unit2 = Entity((16, 24), EntityParams())
+#		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))
+		self.unit_combine = Entity((32, 64), EntityParams(), TrainingParams(learning_rate=0.01, momentum=0.9, do_rao_blackwell=True, num_epochs=1000))
 
-
-	def forward(self, v: Mat):
-		res = np.ndarray(shape=(1,0))
-		index = 0
-		for unit in self.objects():
-			size = unit.shape[0]
-			vp = v[index:index+size]
-			index += size
-			h = unit.forward(vp)
-			res = np.concat((res, h), axis=1)
-		return res
+	def forward2(self, images: Mat, labels: Mat):
+		h = self.unit_image.forward(images)
+		h = np.ndarray.flatten(h)
+		v_combine = np.concat((h, labels), axis=0)
+		h_res = self.unit_combine.forward(v_combine)
+		return h_res
 
 	def reconstruct(self, h: Mat):
-		res = np.ndarray(shape=(1,0))
-		index = 0
-		for unit in self.objects():
-			size = unit.shape[1]
-			hp = np.transpose(np.transpose(h)[index:index+size])
-			index += size
-			v = unit.reconstruct(hp)
-			res = np.concat((res, v), axis=1)
-		return res
+		v_combine = self.unit_combine.reconstruct(h)
+		v_image = self.unit_image.reconstruct(np.transpose(np.transpose(v_combine)[0:self.unit_image.shape[1]]))
+		v_label = np.transpose(np.transpose(v_combine)[self.unit_image.shape[1]:self.unit_image.shape[1]+self.unit_image.shape[1]])
+		return np.ndarray.flatten(v_image), np.ndarray.flatten(v_label)
+
+	def train2(self, images: Mat, labels: Mat):
+		train(self.unit_image, images, Status())
+		h11 = self.unit_image.forward(images)
+		x2 = np.concat((h11, labels), axis=1)
+		train(self.unit_combine, x2, Status())
 
 if __name__ == "__main__":
 	work_dir = "results"
@@ -55,20 +54,37 @@ if __name__ == "__main__":
 	test_batch = read_armadillo(os.path.join(prj_root, f"{prj_name}.test.dat"))
 
 	# Train
-	model.train(train_batch)
+	# Create Labeler
+	enc = Label(20, 16, work_dir)
+
+	# Load
+	enc.fitter.load()
+
+	# Encode test labels
+	train_labels = np.array([1, 2, 3, 4, 5])
+	enc.fit(train_labels)
+
+	encoded, _ = enc.encode(train_labels)
+
+	model.train2(train_batch, encoded)
 
 	# save state
 	model.save()
 
-	fig, axes = plt.subplots(1, len(test_batch), figsize=(12, 3))
-	for index, inp in enumerate(test_batch):
+	fig, axes = plt.subplots(3, len(test_batch), figsize=(12, 3))
+	for index, image in enumerate(test_batch):
 		label_zero = np.zeros(shape=16)
-		query = np.concat((inp, label_zero), axis=0)
-		out_normalized = np.ndarray.flatten(model.reconstruct(model.forward(query)))[0:96*96]
-		img = 2*(out_normalized + 0.5)
+		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 = np.reshape(img, (96, 96))
-		axes[index].imshow(img)
-		axes[index].axis('off')
+		axes[0][index].imshow(img)
+		axes[0][index].axis('off')
+		axes[1][index].imshow(np.reshape(labels_reconst, (4,4)))
+		axes[1][index].axis('off')
+		axes[2][index].imshow(np.reshape(encoded[index, :], (4,4)))
+		axes[2][index].axis('off')
 
 	plt.show()