get_embeddings runtime

Hi guys, we really enjoyed your Nicheformer paper and were excited to try it out. When running it on a dataset of 10,000 cells and 958 genes, it took about 2 hours to generate embeddings. Is this runtime expected? Do you have any tips for speeding it up?
We are using an H100 GPU. We used the get_embeddings() function from the experiments folder in the git repo, as suggested by Alejandro.