diff --git a/src/tests/test_linear.py b/src/tests/test_linear.py
index 22fd60f..5f42dc9 100644
--- a/src/tests/test_linear.py
+++ b/src/tests/test_linear.py
@@ -1,44 +1,64 @@
-import os.path
-
-from rbm.layer import Layer
-from rbm.status import Status
-from rbm.train import train
 from rbm.matrix import Mat, np
-from rbm.entity import EntityParams, TrainingParams
+from rbm.entity import Entity, EntityParams, TrainingParams
+from rbm.model import Model
 
 WORK_DIR = "../../results"
 USE_OPTIMIZER = True
 
+class TestModel(Model):
+	def __init__(self, name: str, work_dir: str = '.', do_gaussian_hidden=False):
+		super().__init__(name, work_dir)
+
+		if do_gaussian_hidden:
+			# Hidden gaussian
+			self.unit1 = Entity((3, 64), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=True),
+								TrainingParams(learning_rate=0.01, momentum=0.9, num_epochs=1000))
+		else:
+			# Hidden binary
+			self.unit1 = Entity((3, 64), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False),
+								TrainingParams(learning_rate=0.001, momentum=0.9, num_epochs=1000))
+
+	def forward(self, x: Mat):
+		x = self.unit1.forward(x)
+		return x
+
+	def reconstruct(self, x: Mat):
+		x = self.unit1.reconstruct(x)
+		return x
+
 def linear():
-	# Create params
-	entity_params = EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False)
-	training_params = TrainingParams(learning_rate=0.0001, do_batch_sample=True, num_epochs=10000, momentum=0.9)
-	entity_params.do_rao_blackwell = False
-	entity_params.num_gibbs_samples = 1
+	work_dir = "results"
+	prj_name = "linear"
+	prj_root = "/home/jens/work/repos/Rbm"
 
 	# Create layer
-	layer = Layer("Layer_0", (3, 1, 0, 32), entity_params, training_params)
+	model = TestModel(prj_name, "results")
 
 	# Init weights
-	layer.init(0.01)
+	model.init(0.1)
 
 	# Load weights (if exists)
-	layer.load(os.path.join(WORK_DIR, "linear_layer0_state.npz"))
+	model.load()
 
 	# Prepare training data
-	training_batch = Mat([[0.5,0.5,1], [0.1,0.9,1.0], [0.2,0.5,0.7], [0.9,0.1,1], [0.5,0.2,0.7]], dtype=np.float64)
+	training_batch = (np.random.rand(150, 3, dtype=np.float64) - 0.5)
+
+	# Normalize training data
+	mean_training_batch = np.reshape(np.repeat(np.mean(training_batch, axis=1), 3, axis=0), training_batch.shape)
+	var_training_batch = np.reshape(np.repeat(np.std(training_batch, axis=1), 3, axis=0), training_batch.shape)
+	training_batch = (training_batch - mean_training_batch) / var_training_batch
 
 	# Train layer
-	train(layer.entity, training_batch, Status())
+	model.train(training_batch)
 
 	# Save weights
-	layer.save(os.path.join(WORK_DIR, "linear_layer0_state.npz"))
+	model.save()
 
 	# Test with test data
 	test_batch = training_batch
 	for pattern in test_batch:
-		h = layer.entity.forward(pattern)
-		v = layer.entity.reconstruct(h)
+		h = model.forward(pattern)
+		v = model.reconstruct(h)
 		print(f"P{pattern} : {v}")
 
 if __name__ == "__main__":
diff --git a/src/tests/test_norbs.py b/src/tests/test_norbs.py
index 6be47fe..058aedd 100644
--- a/src/tests/test_norbs.py
+++ b/src/tests/test_norbs.py
@@ -22,7 +22,7 @@ class TestModel(Model):
 		x = self.unit1.forward(x)
 		return x
 
-	def backward(self, x: Mat):
+	def reconstruct(self, x: Mat):
 		x = self.unit1.reconstruct(x)
 		return x
 
@@ -54,7 +54,7 @@ if __name__ == "__main__":
 
 	fig, axes = plt.subplots(1, len(test_batch), figsize=(12, 3))
 	for index, inp in enumerate(test_batch):
-		out_normalized = model.backward(model.forward(inp))
+		out_normalized = model.reconstruct(model.forward(inp))
 		img = 2*(out_normalized + 0.5)
 		img = np.reshape(img, (96, 96))
 		axes[index].imshow(np.asnumpy(img))