Make the default device for policies "cpu".

HumanCompatibleAI · Dec 11, 2023 · 39647d7 · 39647d7
1 parent 629ef9a
commit 39647d7
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/src/imitation/algorithms/bc.py b/src/imitation/algorithms/bc.py
@@ -249,7 +249,7 @@ def __getstate__(self):
 
 def reconstruct_policy(
     policy_path: str,
-    device: Union[th.device, str] = "auto",
+    device: Union[th.device, str] = "cpu",
 ) -> policies.ActorCriticPolicy:
     """Reconstruct a saved policy.
 
@@ -285,7 +285,7 @@ def __init__(
         optimizer_kwargs: Optional[Mapping[str, Any]] = None,
         ent_weight: float = 1e-3,
         l2_weight: float = 0.0,
-        device: Union[str, th.device] = "auto",
+        device: Union[str, th.device] = "cpu",
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Builds BC.

diff --git a/src/imitation/algorithms/dagger.py b/src/imitation/algorithms/dagger.py
@@ -100,7 +100,7 @@ def reconstruct_trainer(
     scratch_dir: types.AnyPath,
     venv: vec_env.VecEnv,
     custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
-    device: Union[th.device, str] = "auto",
+    device: Union[th.device, str] = "cpu",
 ) -> "DAggerTrainer":
     """Reconstruct trainer from the latest snapshot in some working directory.
 

diff --git a/src/imitation/algorithms/sqil.py b/src/imitation/algorithms/sqil.py
@@ -119,7 +119,7 @@ def __init__(
         observation_space: spaces.Space,
         action_space: spaces.Space,
         demonstrations: algo_base.AnyTransitions,
-        device: Union[th.device, str] = "auto",
+        device: Union[th.device, str] = "cpu",
         n_envs: int = 1,
         optimize_memory_usage: bool = False,
     ):