search

google-research / discs

DISCS: The code base for the Benchmark for Discrete Sampling
Apache License 2.0
47 stars 8 forks source link

VMAP issue in Combinatorial optimization benchmark #63

Open alstn12088 opened 1 year ago

alstn12088 commented 1 year ago

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 commented 1 year ago

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.