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!

首先阅读 sampling_generation_helpers.py 第 58 行附近和 sampling_generation.py 第 94 行附近硬编码的维度断言，然后将它们与 finetuned_replogle.ckpt 的维度以及提供的 sampling 配置进行比较。完成的标准是使用该 checkpoint 可以运行 sampling，并且不存在 2000 个基因的 shape mismatch。