Use PjRt GPU client (#4038)

Commit

2 years ago

Use PjRt GPU client (#4038) * Use PjRt GPU client * Only allow access to GPU device 0 This change also enables CI tests on GPU PjRt