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!

Beginne damit, die fest codierten Dimensionsassertionen in sampling_generation_helpers.py um Zeile 58 und in sampling_generation.py um Zeile 94 zu lesen, und vergleiche sie dann mit den Dimensionen von finetuned_replogle.ckpt und der bereitgestellten Sampling-Konfiguration. Fertig ist es, wenn das Sampling mit dem Checkpoint läuft und keine Shape-Diskrepanz bei 2000 Genen auftritt.