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

Fix: pass device for SAC

Browse source
This commit is contained in:
Antonin RAFFIN 2020-04-06 15:17:30 +02:00
parent c97dbb3173
commit c2c5d0798f
1 changed files with 2 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
torchy_baselines/sac/policies.py
View file

@ -245,7 +245,8 @@ class SACPolicy(BasePolicy):
'features_dim': self.features_dim,
'net_arch': self.net_arch,
'activation_fn': self.activation_fn,
'normalize_images': normalize_images
'normalize_images': normalize_images,
'device': device
}
self.actor_kwargs = self.net_args.copy()
sde_kwargs = {