[RLlib] Discussion 4986: OU Exploration (torch) crashes when restorin…

…g from checkpoint. (ray-project#22245)
simonsays1980 · Feb 27, 2022 · c80684b · c80684b
1 parent f4d3b91
commit c80684b
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/rllib/utils/exploration/ornstein_uhlenbeck_noise.py b/rllib/utils/exploration/ornstein_uhlenbeck_noise.py
@@ -263,10 +263,10 @@ def get_state(self, sess: Optional["tf.Session"] = None):
     def set_state(self, state: dict, sess: Optional["tf.Session"] = None) -> None:
         if self.framework == "tf":
             self.ou_state.load(state["ou_state"], session=sess)
-        elif isinstance(self.ou_state, np.ndarray) or (
-            torch and torch.is_tensor(self.ou_state)
-        ):
+        elif isinstance(self.ou_state, np.ndarray):
             self.ou_state = state["ou_state"]
+        elif torch and torch.is_tensor(self.ou_state):
+            self.ou_state = torch.from_numpy(state["ou_state"])
         else:
             self.ou_state.assign(state["ou_state"])
         super().set_state(state, sess=sess)