diff --git a/src/tests/test_faces_sub_image.py b/src/tests/test_faces_sub_image.py
index b6df614..fcac37a 100644
--- a/src/tests/test_faces_sub_image.py
+++ b/src/tests/test_faces_sub_image.py
@@ -19,12 +19,12 @@ N_IMAGES  = 50
 
 
 class TestModel(Model):
-	def __init__(self, name: str, work_dir: str = '.', l1_lambda: float = 0.0):
+	def __init__(self, name: str, work_dir: str = '.', l1_lambda: float = 0.0, num_epochs: int = 1000):
 		super().__init__(name, work_dir)
 		self.unit1 = Entity(
 			(N_VIS, N_HID),
 			EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False),
-			TrainingParams(learning_rate=0.001, momentum=0.9, num_epochs=1000, mini_batch_size=1000,
+			TrainingParams(learning_rate=0.001, momentum=0.9, num_epochs=num_epochs, mini_batch_size=1000,
 			               l1_lambda=l1_lambda)
 		)
 
@@ -178,6 +178,8 @@ if __name__ == '__main__':
 	                help='Use single-channel grayscale patches (default: false)')
 	ap.add_argument('--l1_lambda',  type=float, default=0.0,
 	                help='L1 regularisation strength (default: 0.0)')
+	ap.add_argument('--num_epochs', type=int,   default=1000,
+	                help='Number of training epochs (default: 1000)')
 	args = ap.parse_args()
 
 	if args.grayscale:
@@ -188,7 +190,7 @@ if __name__ == '__main__':
 	prj_name = 'faces_sub_image_gray' if args.grayscale else 'faces_sub_image'
 	work_dir = 'results'
 
-	model = TestModel(prj_name, work_dir, l1_lambda=args.l1_lambda)
+	model = TestModel(prj_name, work_dir, l1_lambda=args.l1_lambda, num_epochs=args.num_epochs)
 	model.init(0.001)
 
 	if args.load_model: