saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-05-24 22:25:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix: don't change the device when saving

Browse source
This commit is contained in:
Antonin RAFFIN 2020-03-31 18:38:37 +02:00
parent 71ce9ef2f4
commit 6e470d0f72
1 changed files with 0 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
torchy_baselines/common/policies.py
View file

@ -197,11 +197,7 @@ class BasePolicy(nn.Module):
:param path: (str)
"""
previous_device = self.device
# Convert to cpu before saving
self = self.to('cpu')
th.save(self.state_dict(), path)
self = self.to(previous_device)
def load(self, path: str) -> None:
"""
@ -211,7 +207,6 @@ class BasePolicy(nn.Module):
:param path: (str)
"""
self.load_state_dict(th.load(path))
self = self.to(self.device)
def load_from_vector(self, vector: np.ndarray):
"""