From d88898e4772acc2fb45f72b0183612bc4829eb5a Mon Sep 17 00:00:00 2001
From: tastelikefeet <58414341+tastelikefeet@users.noreply.github.com>
Date: Sun, 1 Oct 2023 23:09:09 +0800
Subject: [PATCH] fix work_dir not set in trainer(#573)

* fix work_dir

* fix running
---
 modelscope/trainers/trainer.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/modelscope/trainers/trainer.py b/modelscope/trainers/trainer.py
index a3707918..25f948bc 100644
--- a/modelscope/trainers/trainer.py
+++ b/modelscope/trainers/trainer.py
@@ -181,8 +181,20 @@ class EpochBasedTrainer(BaseTrainer):
                 compile_options = {}
             self.model = compile_model(self.model, **compile_options)
 
-        if 'work_dir' in kwargs:
+        if kwargs.get('work_dir', None) is not None:
             self.work_dir = kwargs['work_dir']
+            if 'train' not in self.cfg:
+                self.cfg['train'] = ConfigDict()
+            self.cfg['train']['work_dir'] = self.work_dir
+            if 'checkpoint' in self.cfg['train']:
+                if 'period' in self.cfg['train']['checkpoint']:
+                    self.cfg['train']['checkpoint']['period'][
+                        'save_dir'] = self.work_dir
+                if 'best' in self.cfg['train']['checkpoint']:
+                    self.cfg['train']['checkpoint']['best'][
+                        'save_dir'] = self.work_dir
+            if 'logging' in self.cfg['train']:
+                self.cfg['train']['logging']['out_dir'] = self.work_dir
         else:
             self.work_dir = self.cfg.train.get('work_dir', './work_dir')