diff --git a/src/rbm/rbm.py b/src/rbm/rbm.py
index 78299f5..e099c9d 100644
--- a/src/rbm/rbm.py
+++ b/src/rbm/rbm.py
@@ -75,7 +75,10 @@ def main(prj_name: str = "test"):
 
 	# Load train data
 	training_data = read_armadillo(os.path.join(prj_root, f"{prj_name}.training.dat"))
-	test_data = read_armadillo(os.path.join(prj_root, f"{prj_name}.test.dat"))
+	try:
+		test_data = read_armadillo(os.path.join(prj_root, f"{prj_name}.test.dat"))
+	except FileNotFoundError:
+		test_data = training_data
 
 	# Prepare status listener
 	my_status = MyStatus(stack, test_data)
@@ -87,7 +90,8 @@ def main(prj_name: str = "test"):
 	stack.state_save()
 
 if __name__ == "__main__":
-	main("norb_small_16h_v2")
+#	main("norb_small_16h_v2")
+	main("mnist_2")
 	cv.destroyAllWindows()
 
 	print("Test: [passed]")