If jax is installed with CUDA, running the CLI will use the GPU, be slow and use 75% VRAM

[LeFiz]

If jax is installed with CUDA, running the CLI will use the GPU, be slow and use 75% VRAM

Runtime as is now (Env running on GPU):

time luxai-s3 Lux-Design-S3/kits/python/main.py Lux-Design-S3/kits/python/main.py 
Init class episode 6.864945411682129s
Agent setup: 0.0012392997741699219s
Env reset: 0.7129645347595215s
First turn: 2.120466709136963s
Finish game: 2.962181806564331s
Cleanup: 0.006840705871582031s
Time Elapsed:  12.669206142425537
Rewards:  {'player_0': array(3, dtype=int32), 'player_1': array(2, dtype=int32)}

real    0m14,492s
user    0m22,287s
sys     0m3,460s

And if running on CPU (via jax.config.update("jax_platform_name", "cpu")):

time luxai-s3 Lux-Design-S3/kits/python/main.py Lux-Design-S3/kits/python/main.py 
Init class episode 3.650555372238159s
Agent setup: 0.0011093616485595703s
Env reset: 0.3220534324645996s
First turn: 1.0807034969329834s
Finish game: 1.4905433654785156s
Cleanup: 0.0054264068603515625s
Time Elapsed:  6.550921678543091
Rewards:  {'player_0': array(1, dtype=int32), 'player_1': array(4, dtype=int32)}

real    0m8,237s
user    0m14,952s
sys     0m1,194s

[StoneT2000]

This is likely due to the compile step for GPUs being much worse maybe. One fix is to force add a flag to turn on/off the use of the GPU backend with default being False since the tool is for one-time episode runs usually.