VMAP issue in Combinatorial optimization benchmark

[alstn12088]

Hello!

The error is the following.

Note I tried multiple versions of dependencies and multiple people tried the same things.

new_x, state, acc = stp_burnin(
ValueError: vmap got inconsistent sizes for array axes to be mapped:
  * most axes (8 of them) had size 1024, e.g. axis 0 of argument rng of type uint32[1,2];
  * some axes (6 of them) had size 1, e.g. axis 0 of argument model_param['edge_weight'] of type float32[1024,64]

[hyeonahkimm]

I also have encountered the same issue when I use multiple GPUs (even for a single GPU, there is an error with a different message).

I found the configuration of num_models is duplicated (in the experiment and model config).
I've resolved this issue by making the two configurations the same.

1. add num_models in experiment config (I guess the default value is 1)
2. change num_models to 1 in model config

It works regardless of the number of GPUs. Good luck.