From 9e3f035fa71ef15c6dcf10a4388630e6c634dc40 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=A1=8C=E5=97=94?= Date: Tue, 25 Oct 2022 10:13:48 +0800 Subject: [PATCH] fix a ut bug --- tests/trainers/test_ofa_trainer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/trainers/test_ofa_trainer.py b/tests/trainers/test_ofa_trainer.py index 786599bb..46dc5c8b 100644 --- a/tests/trainers/test_ofa_trainer.py +++ b/tests/trainers/test_ofa_trainer.py @@ -79,7 +79,7 @@ class TestOfaTrainer(unittest.TestCase): with open(config_file, 'w') as writer: json.dump(self.finetune_cfg, writer) - pretrained_model = 'damo/ofa_image-caption_coco_large_en' + pretrained_model = 'damo/ofa_image-caption_coco_distilled_en' args = dict( model=pretrained_model, work_dir=WORKSPACE, @@ -97,8 +97,8 @@ class TestOfaTrainer(unittest.TestCase): trainer.train() self.assertIn(ModelFile.TORCH_MODEL_BIN_FILE, - os.path.join(WORKSPACE, 'output')) - shutil.rmtree(WORKSPACE) + os.listdir(os.path.join(WORKSPACE, 'output'))) + # shutil.rmtree(WORKSPACE) if __name__ == '__main__':