Wrong size of 'x' in each batch

[ekernf01]

Hi scFormer devs, I am seeing an error when testing out the scFormer perturbation prediction demo.

max num of non_zeros_gene_idx: 5012
min num of non_zeros_gene_idx: 4652
mean num of non_zeros_gene_idx: 4944.631578947368
Traceback (most recent call last):
  File "/home/ekernf01/Desktop/jhu/research/projects/perturbation_prediction/cell_type_knowledge_transfer/test_scformer.py", line 45, in <module>
    pt.simple_training()
  File "/home/ekernf01/Desktop/jhu/research/projects/perturbation_prediction/cell_type_knowledge_transfer/scformer_perturb/scformer_perturb.py", line 167, in simple_training
    self.train(
  File "/home/ekernf01/Desktop/jhu/research/projects/perturbation_prediction/cell_type_knowledge_transfer/scformer_perturb/scformer_perturb.py", line 241, in train
    pert_flags = x[:, 1].long().view(batch_size, self.n_genes)
IndexError: index 1 is out of bounds for dimension 1 with size 1

Here are the details on a batch causing this error. It looks like x is too small along axis 1. Can you tell from this what needs to be fixed? Thank you!

DataBatch(x=[161440, 1], y=[32, 5045], pert_idx=[32], de_idx=[32], pert=[32], batch=[161440], ptr=[33])