diff --git a/src/rbm/matrix.py b/src/rbm/matrix.py
index 441cda8..f47a12a 100644
--- a/src/rbm/matrix.py
+++ b/src/rbm/matrix.py
@@ -1,7 +1,6 @@
-import time
 from typing import TypeAlias
 
-USE_CUDA = 1
+USE_CUDA = 0
 if USE_CUDA:
 	import cupy as np
 	Mat: TypeAlias = np.array
@@ -36,3 +35,22 @@ def rms_error_accu(d_err: Mat):
 	s =  np.sum(d_err_squared) / d_err.size
 	return s
 
+def read_armadillo(filename: str) -> Mat:
+	result = None
+	with open(filename) as fp:
+		identifier = fp.readline().replace("\n", '')
+		if "ARMA_MAT_TXT_FN008" not in identifier:
+			raise Exception("Not a armadillo data file!")
+
+		line = fp.readline().replace("\n", '').split(' ')
+		shape = [int(s) for s in line]
+		print(f"shape: {shape}")
+
+		result = np.zeros(shape=shape, dtype=np.float64)
+		for row in range(shape[0]):
+			line = fp.readline().replace("\n", '').split('  ')
+			line = line[1:]
+			data = [float(s) for s in line]
+			result[row, :] = Mat(data)
+
+	return result
diff --git a/src/tests/__init__.py b/src/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/tests/test_norbs.py b/src/tests/test_norbs.py
new file mode 100644
index 0000000..93ea616
--- /dev/null
+++ b/src/tests/test_norbs.py
@@ -0,0 +1,60 @@
+import os
+
+import matplotlib.pyplot as plt
+
+from rbm.model import Model
+from rbm.entity import Entity, EntityParams
+from rbm.matrix import Mat, np, read_armadillo
+from rbm.train import TrainingParams
+
+import matplotlib.pyplot as plot
+
+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))
+
+	def forward(self, x: Mat):
+		x = self.unit1.forward(x)
+		return x
+
+	def backward(self, x: Mat):
+		x = self.unit1.reconstruct(x)
+		return x
+
+if __name__ == "__main__":
+	work_dir = "results"
+	prj_name = "norb_small_16h_v2"
+	prj_root = "/home/jens/work/repos/Rbm"
+
+	# Create model
+	model = TestModel("norb_small_16h_v2", "results")
+
+	# Init state
+	model.init(0.01)
+
+	# load state
+	model.load()
+
+	# Load train data
+	batch = read_armadillo(os.path.join(prj_root, f"{prj_name}.training.dat"))
+
+	# Train
+	model.train(batch, TrainingParams(learning_rate=0.00001, momentum=0.9, do_rao_blackwell=True, num_epochs=100, num_gibbs_samples=3))
+
+	# save state
+	model.save()
+
+	num_patterns = len(batch)
+	fig, axes = plt.subplots(1, num_patterns, figsize=(12, 3))
+	for index, inp in enumerate(batch):
+		out_normalized = model.backward(model.forward(inp))
+		img = 2*(out_normalized + 0.5)
+		img = np.reshape(img, (96, 96))
+		axes[index].imshow(img)
+		axes[index].axis('off')
+
+	plt.show()
+
+
+
diff --git a/src/tests/test_rbm.py b/src/tests/test_rbm.py
index 133de36..d6bd2bf 100644
--- a/src/tests/test_rbm.py
+++ b/src/tests/test_rbm.py
@@ -4,7 +4,7 @@ from argparse import ArgumentParser
 from rbm.stack_factory import StackFactory
 from rbm.status import Status
 from rbm.stack_deep import StackDeep
-from rbm.matrix import Mat, np, convert
+from rbm.matrix import Mat, np, convert, read_armadillo
 
 def cv_show(name: str, vec: Mat, shape):
 	img = cv.Mat(convert(np.resize(vec, shape)))
@@ -41,25 +41,6 @@ class MyStatus(Status):
 
 		return do_continue
 
-def read_armadillo(filename: str) -> Mat:
-	result = None
-	with open(filename) as fp:
-		identifier = fp.readline().replace("\n", '')
-		if "ARMA_MAT_TXT_FN008" not in identifier:
-			raise Exception("Not a armadillo data file!")
-
-		line = fp.readline().replace("\n", '').split(' ')
-		shape = [int(s) for s in line]
-		print(f"shape: {shape}")
-
-		result = np.zeros(shape=shape, dtype=np.float64)
-		for row in range(shape[0]):
-			line = fp.readline().replace("\n", '').split('  ')
-			line = line[1:]
-			data = [float(s) for s in line]
-			result[row, :] = Mat(data)
-
-	return result
 
 def main(prj_name: str = "test"):
 	work_dir = "../../results"