update save and load
This commit is contained in:
@@ -93,14 +93,14 @@ class ACAgent(Agent):
|
||||
def diagnostics(self):
|
||||
return {'value': self.last_state_value}
|
||||
|
||||
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)
|
||||
|
||||
# def load_ac_agent(path, encoder):
|
||||
|
||||
Reference in New Issue
Block a user