Update README.md

README.md

@@ -40,9 +40,12 @@ Here is how to use the model to extract features from the pre-trained backbone:
 ```python
 import torch
 from transformers import AutoModel
-model = AutoModel.from_pretrained("vector-institute/atomformer-base", trust_remote_code=True)
+model = AutoModel.from_pretrained("vector-institute/atomformer-base",
+                                   trust_remote_code=True)
 
-input_ids, coords, attn_mask = torch.randint(0, 100, (1, 10)), torch.randn(1, 10, 3), torch.ones(1, 10)
+input_ids = torch.randint(0, 50, (1, 10))
+coords = torch.randn(1, 10, 3)
+attn_mask = torch.ones(1, 10)
 
 output = model(input_ids, coords=coords, attention_mask=attention_mask)
 output[0].shape #  (torch.Size([1, 10, 768])