Perturbation tutorial not working

[yinjun111]

Hi

Thanks for sharing the model and the tutorial with the community. It is a great resource to learn from this foundation model and use to interpret our data.

I tried to test the Tutorial_Perturbation.ipynb using the example data set "adamson". The script works fine till the "train()" step. It reports:

Cell In[133], line 22
20 x: torch.Tensor = batch_data.x # (batch_size * n_genes, 2)
21 ori_gene_values = x[:, 0].view(batch_size, n_genes)
---> 22 pert_flags = x[:, 1].long().view(batch_size, n_genes)
23 target_gene_values = batch_data.y # (batch_size, n_genes)
25 if include_zero_gene in ["all", "batch-wise"]:

IndexError: index 1 is out of bounds for dimension 1 with size 1

The error seems to come from the pert_flags loading from batch_data.x, that batch_data.x is a one-column matrix, it doesn't have the 2nd column as required by the code. I also tried "norman" data set, and it gives the same error.

Please help to correct the tutorial. Thanks in advance.