diff --git a/inference.py b/inference.py
index 6054b01..8d0a572 100644
--- a/inference.py
+++ b/inference.py
@@ -28,7 +28,7 @@ Tool to load model and predict for given image.
 """
 
 class sbb_predict:
-    def __init__(self,image, model, task, config_params_model, patches, save, ground_truth, xml_file, out, min_area):
+    def __init__(self,image, model, task, config_params_model, patches, save, save_layout, ground_truth, xml_file, out, min_area):
         self.image=image
         self.patches=patches
         self.save=save