fix train

dc7d4b00 · LDOUBLEV · c4232b8a · dc7d4b00
Commit dc7d4b00 authored Dec 18, 2021 by LDOUBLEV
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 4 deletions

tools/program.py tools/program.py +2 -4

No files found.
--- a/tools/program.py
+++ b/tools/program.py
@@ -227,10 +227,6 @@ def train(config,
            images = batch[0]
            if use_srn:
                model_average = True
-            if model_type == 'table' or extra_input:
-                preds = model(images, data=batch[1:])
-            if model_type == "kie":
-                preds = model(batch)

            train_start = time.time()
            # use amp
@@ -243,6 +239,8 @@ def train(config,
            else:
                if model_type == 'table' or extra_input:
                    preds = model(images, data=batch[1:])
+                if model_type == "kie":
+                    preds = model(batch)
                else:
                    preds = model(images)
            loss = loss_class(preds, batch)