saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-05-18 21:30:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix device

Browse source
This commit is contained in:
Antonin Raffin 2019-11-13 14:38:18 +01:00
parent da325a0ba7
commit 8aac10f3fa
1 changed files with 1 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
torchy_baselines/td3/policies.py
View file

@ -40,7 +40,7 @@ class Actor(BaseNetwork):
if self.full_std:
return self.log_std
# Reduce the number of parameters:
return th.ones((self.latent_dim, self.action_dim)) * self.log_std
return th.ones((self.latent_dim, self.action_dim)).to(self.log_std.device) * self.log_std
def get_distribution_stats(self, obs, action):
with th.no_grad():