Error in training very deep GNN for PCQM4M-LSC

**Hi, All**

I am encounter a problem after running the following training command of `ogb_lsc/pcq/` at https://github.com/deepmind/deepmind-research/tree/master/ogb_lsc/pcq#retraining-our-model
```
/bin/bash run_training.sh -r ${HOME}/pcq/
```
The error is
```
I0126 23:29:54.918133 139951828096768 utils.py:234] [jaxline] experiment init starting...
E0126 23:29:54.918646 139951828096768 utils.py:239] [jaxline] experiment init failed with error.
RuntimeError: Visible devices cannot be modified after being initialized
```
I have tried both small and large model by setting the `debug=True/False` in `config.py`, both of them get the same error.
I currently only have CPU and GPU on my computer.