diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index c262ea308d..51335940a8 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -147,7 +147,7 @@ def _allocate_gpus_via_placements(self, worker_placements, num_gpus_per_actor) - cvd = ",".join(str(g) for g in wp.gpu_ids) env_vars = dict(env_vars_base) env_vars["CUDA_VISIBLE_DEVICES"] = cvd - TrainRayActor = ray.remote(num_gpus=1, runtime_env={"env_vars": env_vars})(actor_impl) + TrainRayActor = ray.remote(runtime_env={"env_vars": env_vars})(actor_impl) actor = TrainRayActor.options( num_cpus=num_gpus_per_actor, num_gpus=num_gpus_per_actor,