saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-09-16 22:20:26 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix initializing CUDA even when device="cpu" is used. (#194)

Browse source 
* Fall back to 'cpu' device in policies instead of 'auto'

* Update changelog
This commit is contained in:
Anssi 2020-10-18 21:51:56 +03:00 committed by GitHub
parent 97b81f9e9e
commit 37f48aa979
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 6 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
docs/misc/changelog.rst
View file

@ -19,6 +19,7 @@ Bug Fixes:
- Fix GAE computation for on-policy algorithms (off-by one for the last value) (thanks @Wovchena) - Fix GAE computation for on-policy algorithms (off-by one for the last value) (thanks @Wovchena)
- Fix ignoring the exclude parameter when recording logs using json, csv or log as logging format (@SwamyDev) - Fix ignoring the exclude parameter when recording logs using json, csv or log as logging format (@SwamyDev)
- Make ``make_vec_env`` support the ``env_kwargs`` argument when using an env ID str (@ManifoldFR) - Make ``make_vec_env`` support the ``env_kwargs`` argument when using an env ID str (@ManifoldFR)
- Fix model creation initializing CUDA even when `device="cpu"` is provided
Deprecations: Deprecations:
^^^^^^^^^^^^^ ^^^^^^^^^^^^^

8
stable_baselines3/common/policies.py
View file

@ -112,12 +112,12 @@ class BaseModel(nn.Module, ABC):
@property @property
def device(self) -> th.device: def device(self) -> th.device:
"""Infer which device this policy lives on by inspecting its parameters. """Infer which device this policy lives on by inspecting its parameters.
If it has no parameters, the 'auto' device is used as a fallback. If it has no parameters, the 'cpu' device is used as a fallback.
:return:""" :return:"""
for param in self.parameters(): for param in self.parameters():
return param.device return param.device
return get_device("auto") return get_device("cpu")
def save(self, path: str) -> None: def save(self, path: str) -> None:
""" """
@ -447,7 +447,9 @@ class ActorCriticPolicy(BasePolicy):
# Note: If net_arch is None and some features extractor is used, # Note: If net_arch is None and some features extractor is used,
# net_arch here is an empty list and mlp_extractor does not # net_arch here is an empty list and mlp_extractor does not
# really contain any layers (acts like an identity module). # really contain any layers (acts like an identity module).
self.mlp_extractor = MlpExtractor(self.features_dim, net_arch=self.net_arch, activation_fn=self.activation_fn) self.mlp_extractor = MlpExtractor(
self.features_dim, net_arch=self.net_arch, activation_fn=self.activation_fn, device=self.device
)
def _build(self, lr_schedule: Callable[[float], float]) -> None: def _build(self, lr_schedule: Callable[[float], float]) -> None:
""" """