From 7aeb60c20c3b74a50d79fe5708c73518ea5d330c Mon Sep 17 00:00:00 2001
From: liuhongen1234567 <2998388548@qq.com>
Date: Tue, 21 Jan 2025 13:09:47 +0000
Subject: [PATCH] repair train bug in multi gpu

---
 configs/rec/rec_latex_ocr.yml | 1 +
 tools/train.py                | 5 ++++-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/configs/rec/rec_latex_ocr.yml b/configs/rec/rec_latex_ocr.yml
index b6650686b6..377e0f1f66 100644
--- a/configs/rec/rec_latex_ocr.yml
+++ b/configs/rec/rec_latex_ocr.yml
@@ -19,6 +19,7 @@ Global:
   rec_char_dict_path:  ppocr/utils/dict/latex_ocr_tokenizer.json
   save_res_path: ./output/rec/predicts_latexocr.txt
   d2s_train_image_shape: [1,256,256]
+  find_unused_parameters: True
 
 Optimizer:
   name: AdamW
diff --git a/tools/train.py b/tools/train.py
index 0a2e2e6458..7ce0faf03a 100755
--- a/tools/train.py
+++ b/tools/train.py
@@ -217,7 +217,10 @@ def main(config, device, logger, vdl_writer, seed):
     )
 
     if config["Global"]["distributed"]:
-        model = paddle.DataParallel(model)
+        find_unused_parameters = config["Global"].get("find_unused_parameters", False)
+        model = paddle.DataParallel(
+            model, find_unused_parameters=find_unused_parameters
+        )
     # start train
     program.train(
         config,