Stable Baselines - model.predict: how to select the GPU

Question

I'm using Stable-Baselines3 and I'd like to test a model I've just trained.

How can I assign a specific GPU to the execution of the model.predict method?

Answer 1

You need to assign the device to the model when you create it, then predict would automatically use the chosen device.

E.g.

model = DQN(..., device='cuda:0', ...)
model.predict(...)    # this would use the cuda:0 device