How to enable custom Gym environment access GPUs?

How severe does this issue affect your experience of using Ray?

  • Medium: It contributes to significant difficulty in completing my task, but I can work around it.

Hi, I have a custom Gym environment defined which needs access to GPUs as the reward function is estimated by a pre-trained model. How can I go about doing this in RLlib? I think this is a pretty common need but I’ve not been able to find any resource to do this.

I’ve tried to use @ray.remote but it throws up errors. Any help is appreciated.

Thanks!