How to specify the GPU that loads the model weights when applying to IsaacGym

Denys88 / rl_games

RL implementations

MIT License

863 stars 146 forks source link

Open Charlie0257 opened 9 months ago

Charlie0257 commented 9 months ago

Since I trained the corresponding model weights on the server, how do I specify the corresponding loaded GPU ID on my own single GPU machine?

Denys88 commented 8 months ago

you can setup device in yaml config. But by default it is just 'cuda' so it should choose your only device.

Charlie0257 commented 8 months ago

@Denys88 Thanks for your reply!

Setting up device in yaml config does not work!

The reason is return safe_filesystem_op(torch.load, filename) in torch_ext.py don't specify device!

Do you have any suggestions?