hongshaorou/YOPO
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix file name bug when saving

Browse Source
This commit is contained in:
TJU_Lu 2025-06-21 15:30:47 +08:00
parent 4e09f3b64a
commit e019bd65de
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
YOPO/policy/yopo_trainer.py
View File

@ -72,7 +72,7 @@ class YopoTrainer:
self.eval_one_epoch(self.epoch_i) self.eval_one_epoch(self.epoch_i)
if save_interval is not None and (self.epoch_i + 1) % save_interval == 0: if save_interval is not None and (self.epoch_i + 1) % save_interval == 0:
self.progress_log.console.log("Saving model...") self.progress_log.console.log("Saving model...")
policy_path = self.tensorboard_path + "/epoch{}.pth".format(epoch + 1, 0) policy_path = self.tensorboard_path + "/epoch{}.pth".format(self.epoch_i + 1, 0)
torch.save(self.policy.state_dict(), policy_path) torch.save(self.policy.state_dict(), policy_path)
self.progress_log.console.log("Train YOPO Finish!") self.progress_log.console.log("Train YOPO Finish!")
self.progress_log.remove_task(total_progress) self.progress_log.remove_task(total_progress)