Update README.md

README.md

@@ -23,7 +23,13 @@ To handle missing sections, we employ special tokens.
 We also utilise an attention mask with non-causal masking for the image embeddings and a causal mask for the report token embeddings.
 
 ## How to use:
-```
+
+```python
+import torch
+from torchvision.transforms import v2
+import transformers
+
+
 tokenizer = transformers.AutoTokenizer.from_pretrained('aehrc/cxrmate-rrg24')
 model = transformers.AutoModel.from_pretrained('aehrc/cxrmate-rrg24', trust_remote_code=True)
 transforms = v2.Compose(
@@ -38,7 +44,7 @@ transforms = v2.Compose(
 )
 image = transforms(image)
 output_ids = model.generate(
-    pixel_values=image.unsqueeze(0).unsqueeze(0),
+    pixel_values=images,
     max_length=512,
     bad_words_ids=[[tokenizer.convert_tokens_to_ids('[NF]')], [tokenizer.convert_tokens_to_ids('[NI]')]],
     num_beams=4,