From 37f48aa9794dd65d4814a13378e1a39458e3b16e Mon Sep 17 00:00:00 2001 From: Anssi Date: Sun, 18 Oct 2020 21:51:56 +0300 Subject: [PATCH] Fix initializing CUDA even when `device="cpu"` is used. (#194) * Fall back to 'cpu' device in policies instead of 'auto' * Update changelog --- docs/misc/changelog.rst | 1 + stable_baselines3/common/policies.py | 8 +++++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 9181878..3f6e087 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -19,6 +19,7 @@ Bug Fixes: - 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) - 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: ^^^^^^^^^^^^^ diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index 8a0fbbb..fa21018 100644 --- a/stable_baselines3/common/policies.py +++ b/stable_baselines3/common/policies.py @@ -112,12 +112,12 @@ class BaseModel(nn.Module, ABC): @property def device(self) -> th.device: """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:""" for param in self.parameters(): return param.device - return get_device("auto") + return get_device("cpu") 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, # net_arch here is an empty list and mlp_extractor does not # 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: """