QDax
QDax copied to clipboard
Filling Genotype with jnp.zeros doesn't allocate GPU memory
Hello :),
When running more memory intensive QD algorithms, it is possible to be out of memory when filling up the repertoire. I found a comment on the jax repo, it seems that jnp.zeros is not really allocating memory which means that every time we add individuals, we start allocating memory. I haven't specifically tested the memory usage but when changing the initialization of the genotype to jnp.nan, it seems that I get less OOM errors (might be an illusion). So we might want to investigate this a little further.