How to use multiple GPUs?

[cenwangumass]

I have multiple GPUs on a machine and am trying to understand how to use them. In the C++ binding function execute I can see there is a line PjRtDevice *device = client->devices()[0]; which seems to always use the first available GPU. I tried to test changing 0 to 1, 2, etc., but am having the following error:

Execution of replica 0 failed: INVALID_ARGUMENT: Buffer passed to Execute() as argument 0 to replica 0 is on device gpu:3, but replica is assigned to device gpu:0.

Is there any suggestions how I should modified the code to use multiple GPUs? Thanks.