Improved policy executer

2024-01-16 23:22:05 +00:00
parent d1074281c4
commit b87ad1bf42
7 changed files with 1328 additions and 101 deletions
--- a/src/trainers/trainer.py
+++ b/src/trainers/trainer.py
@@ -249,7 +249,7 @@ class Trainer:

    def finish_training(self, task):
        if self.best_score is not None:
-            self.model.load_state_dict(torch.load("checkpoint.pt"))
+            self.model = torch.load("checkpoint.pt")
        self.model.eval()