Force worker to use CPU for model and GPU for env

RLlib
1

During evaluation, I want my environment to have access to the GPU to render the episode with fancy graphics, but the worker should run the model only on CPU.

The reason I want the worker to run on CPU is that the trainer and validation environment is already in the GPU memory (an there isn’t much memory left).

evaluation_config = {
                "num_gpus_per_worker": 0.1
            }

To use 1 GPU for training and evaluation environment, I set the num_gpus_per_worker to < 1.0 to share it. But this results in validation worker seeing, it has access to GPU and allocates ray::RolloutWorker 566MiB on GPU memory.

Any tips on forcing worker to use CPU model even though it has access to GPU?

Ray version: 2.2.0
Platform: TF

Thanks!