Commit 6af5a54c authored by Julien Chaumond's avatar Julien Chaumond
Browse files

[Trainer] reuse constant

parent 7c2a32ff
...@@ -358,7 +358,7 @@ class Trainer: ...@@ -358,7 +358,7 @@ class Trainer:
else: else:
assert model is self.model assert model is self.model
# Save model checkpoint # Save model checkpoint
output_dir = os.path.join(self.args.output_dir, f"checkpoint-{global_step}") output_dir = os.path.join(self.args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{global_step}")
self.save_model(output_dir) self.save_model(output_dir)
self._rotate_checkpoints() self._rotate_checkpoints()
torch.save(optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt")) torch.save(optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment