merge

Signed-off-by: sven1977 <svenmika1977@gmail.com>
ray-project · sven1977 · Jan 7, 2025 · Dec 23, 2024 · Dec 27, 2024 · Dec 27, 2024
commit f19b34022842187b5c834fc3940dc8446fb3999d
diff --git a/python/ray/data/_internal/planner/plan_udf_map_op.py b/python/ray/data/_internal/planner/plan_udf_map_op.py
@@ -48,7 +48,7 @@
 )
 from ray.data.context import DataContext
 from ray.data.exceptions import UserCodeException
-#from ray.util.rpdb import _is_ray_debugger_post_mortem_enabled
+from ray.util.rpdb import _is_ray_debugger_post_mortem_enabled
 
 
 class _MapActorContext:

diff --git a/python/ray/data/exceptions.py b/python/ray/data/exceptions.py
@@ -6,7 +6,7 @@
 from ray.exceptions import UserCodeException
 from ray.util import log_once
 from ray.util.annotations import DeveloperAPI
-#from ray.util.rpdb import _is_ray_debugger_post_mortem_enabled
+from ray.util.rpdb import _is_ray_debugger_post_mortem_enabled
 
 logger = logging.getLogger(__name__)
 

diff --git a/rllib/BUILD b/rllib/BUILD
@@ -1594,6 +1594,7 @@ py_test(
     name = "utils/tests/test_checkpointable",
     tags = ["team:rllib", "utils"],
     size = "large",
+    data = glob(["utils/tests/old_checkpoints/**"]),
     srcs = ["utils/tests/test_checkpointable.py"]
 )
 

diff --git a/rllib/core/learner/torch/torch_learner.py b/rllib/core/learner/torch/torch_learner.py
@@ -338,7 +338,7 @@ def _set_optimizer_state(self, state: StateDict) -> None:
                 )
             if name in self._named_optimizers:
                 self._named_optimizers[name].load_state_dict(
-                    convert_to_torch_tensor(state_dict["state"])
+                    convert_to_torch_tensor(state_dict["state"], device=self._device)
                 )
 
     @override(Learner)