[Trainer] reuse constant

This commit is contained in:
Julien Chaumond
2020-04-23 11:02:05 -04:00
parent 7c2a32ff88
commit 6af5a54c28

View File

@@ -358,7 +358,7 @@ class Trainer:
else:
assert model is self.model
# 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._rotate_checkpoints()
torch.save(optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))