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!

Inizia leggendo le asserzioni sulle dimensioni codificate direttamente in sampling_generation_helpers.py intorno alla riga 58 e in sampling_generation.py intorno alla riga 94, quindi confrontale con le dimensioni di finetuned_replogle.ckpt e con la configurazione di sampling fornita. Il lavoro è completato quando il sampling viene eseguito con il checkpoint e non si verifica alcuna discrepanza di forma per 2000 geni.