稳定基线 - model.predict：如何选择 GPU

Question

我正在使用 Stable-Baselines3，我想测试我刚刚训练的模型。

如何分配特定的 GPU 来执行

model.predict

方法？

Answer 1

创建模型时需要将设备分配给模型，然后

predict

会自动使用所选的设备。

例如

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