Only convert to numpy if needed

DLR-RM · Jul 12, 2024 · fb1a9f7 · fb1a9f7
1 parent bee9cbe
commit fb1a9f7
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/stable_baselines3/common/prioritized_replay_buffer.py b/stable_baselines3/common/prioritized_replay_buffer.py
@@ -247,7 +247,8 @@ def update_priorities(self, leaf_nodes_indices: np.ndarray, td_errors: th.Tensor
         """
         # Update beta schedule
         self._current_progress_remaining = progress_remaining
-        td_errors = td_errors.detach().cpu().numpy().flatten()
+        if isinstance(td_errors, th.Tensor):
+            td_errors = td_errors.detach().cpu().numpy().flatten()
 
         for leaf_node_idx, td_error in zip(leaf_nodes_indices, td_errors):
             # Proportional prioritization priority = (abs(td_error) + eps) ^ alpha