Error in training very deep GNN for PCQM4M-LSC

[YukeWang96]

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.