RuntimeError: mat1 and mat2 shapes cannot be multiplied during sampling with finetuned_replogle.ckpt

Hello! I'm trying to run the sampling script using the provided finetuned_replogle.ckpt.

It seems this checkpoint was trained on 12626 genes, but script defaults to a 2000-gene space, causing a shape mismatch error:
RuntimeError: mat1 and mat2 shapes cannot be multiplied (128x4000 and 25252x512)

I tried changing the config to match 12626 genes:

data.pad_length=12626
model.hidden_num=[12626,512]
model.input_dim=12626
data.embed_key=X

However, the script still fails because 2000 is strictly hardcoded in several assert statements across the sampling codebase (e.g., sampling_generation_helpers.py line 58 and sampling_generation.py line 94).

How can I correctly run sampling for this finetuned checkpoint? Are there plans to make the gene dimension dynamic instead of hardcoded?

Thanks!

Comienza leyendo las aserciones de dimensiones codificadas de forma fija en sampling_generation_helpers.py alrededor de la línea 58 y en sampling_generation.py alrededor de la línea 94; después, compáralas con las dimensiones de finetuned_replogle.ckpt y con la configuración de sampling proporcionada. Se considera terminado cuando el sampling se ejecuta con el checkpoint y no hay ninguna discrepancia de forma para 2000 genes.