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!

Bắt đầu bằng cách đọc các assertion về kích thước được hardcode trong sampling_generation_helpers.py quanh dòng 58 và sampling_generation.py quanh dòng 94, sau đó so sánh chúng với các kích thước của finetuned_replogle.ckpt và cấu hình sampling được cung cấp. Được xem là hoàn thành khi sampling chạy được với checkpoint và không có lỗi không khớp shape với 2000 gene.