Update README.md

README.md

@@ -40,8 +40,8 @@ model = AutoModel.from_pretrained(
     trust_remote_code=True,
 )
 
-n_genes = jax_config.n_genes
-dummy_gene_expressions = torch.randint(0, jax_config.n_expressions_bins, (1, n_genes))
+n_genes = model.config.n_genes
+dummy_gene_expressions = torch.randint(0, model.config.n_expressions_bins, (1, n_genes))
 torch_output = model(dummy_gene_expressions)
 ```