Assign self.device = device

AlignmentResearch · Oct 11, 2023 · 1978bc6 · 1978bc6
1 parent d800947
commit 1978bc6
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/stable_baselines3/common/recurrent/buffers.py b/stable_baselines3/common/recurrent/buffers.py
@@ -17,6 +17,7 @@
     RecurrentRolloutBufferData,
     RecurrentRolloutBufferSamples,
 )
+from stable_baselines3.common.utils import get_device
 from stable_baselines3.common.vec_env import VecNormalize
 
 
@@ -173,7 +174,7 @@ def __init__(
         self.gamma = gamma
 
         batch_shape = (self.buffer_size, self.n_envs)
-        device = self.device
+        self.device = device = get_device(device)
 
         self.observation_space_example = space_to_example((), observation_space)