From 3ac03e448b7673604eb86f619b27521fca55f34d Mon Sep 17 00:00:00 2001
From: zhifu gao <zhifu.gzf@alibaba-inc.com>
Date: 星期四, 21 三月 2024 01:36:39 +0800
Subject: [PATCH] train & finetune llm-asr (#1519)
---
funasr/train_utils/trainer.py | 6 ++++--
1 files changed, 4 insertions(+), 2 deletions(-)
diff --git a/funasr/train_utils/trainer.py b/funasr/train_utils/trainer.py
index a00b3de..aae4513 100644
--- a/funasr/train_utils/trainer.py
+++ b/funasr/train_utils/trainer.py
@@ -103,6 +103,7 @@
os.makedirs(os.path.join(self.output_dir, "tensorboard"), exist_ok=True)
self.writer = SummaryWriter(os.path.join(self.output_dir, "tensorboard")) if rank == 0 else None
+
def _save_checkpoint(self, epoch, step=None):
@@ -146,7 +147,7 @@
"""
ckpt = os.path.join(resume_path, "model.pt")
if os.path.isfile(ckpt):
- checkpoint = torch.load(ckpt)
+ checkpoint = torch.load(ckpt, map_location="cpu")
self.start_epoch = checkpoint['epoch'] + 1
# self.model.load_state_dict(checkpoint['state_dict'])
src_state = checkpoint['state_dict']
@@ -169,7 +170,8 @@
print(f"Checkpoint loaded successfully from '{ckpt}'")
else:
print(f"No checkpoint found at '{ckpt}', does not resume status!")
-
+
+ self.model.to(self.device)
if self.use_ddp or self.use_fsdp:
dist.barrier()
--
Gitblit v1.9.1