search

PufferAI / PufferLib

Simplifying reinforcement learning for complex game environments
https://pufferai.github.io/
MIT License
1.23k stars 58 forks source link

Allow device to be like cuda:1 #95

Open kywch opened 2 months ago

EliaPiccoli commented 1 month ago

I think you also need to update the torch.cuda.syncronize() calls otherwise it synchs torch.cuda.current_device that - if not set - is not the actual device used. You can use torch.cuda.synchronize(config.device) instead to synch the correct device.

kywch commented 1 month ago

@EliaPiccoli Thank you so much! I followed your suggestion.

EliaPiccoli commented 1 month ago

Hi, me again :D While performing experiments I also noticed that the GPU usage should be adapted while using gpu different from 'cuda:0'. I believe you need to update lines 501-502 adding the gpu device as argument. In order to have the gpu information though you need to add an argument to the Utilization class, which I think should not be a problem right?

kywch commented 1 month ago

Hi Elia, nice catch and thanks! As you suggested, I passed in the device info, and now it should query the correct device.