update save and load

2023-05-31 16:16:24 +08:00
parent 41c70802fd
commit fe2c327f4e
12 changed files with 252 additions and 44 deletions
@@ -38,12 +38,12 @@ class DeepLearningAgent(Agent):
                return goboard.Move.play(point)
        return goboard.Move.pass_turn()

-    def save(self, path):
-        torch.save(self.model.state_dict(), path)
+    def save(self, file_path):
+        torch.save(self.model.state_dict(), file_path)

    @classmethod
-    def load(cls, path, encoder):
+    def load(cls, file_path, encoder):
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        model = AlphaGoModel(encoder.get_input_shape(), is_policy_net=True).to(device)
-        model.load_state_dict(torch.load(path))
+        model.load_state_dict(torch.load(file_path))
        return cls(model, encoder)