diff --git a/tools/program.py b/tools/program.py
index 538d0f62fd..84dc6c1d02 100755
--- a/tools/program.py
+++ b/tools/program.py
@@ -239,7 +239,7 @@ def train(config,
             else:
                 if model_type == 'table' or extra_input:
                     preds = model(images, data=batch[1:])
-                if model_type == "kie":
+                elif model_type == "kie":
                     preds = model(batch)
                 else:
                     preds = model(images)