Update README.md

README.md

@@ -37,8 +37,9 @@ from torch.utils.data import DataLoader
 
 
 mbatch_size = 1
+device = 'cuda'
 
-tokenizer = transformers.AutoTokenizer.from_pretrained('aehrc/cxrmate-rrg24')
+tokenizer = transformers.AutoTokenizer.from_pretrained('aehrc/cxrmate-rrg24').to(device=device)
 model = transformers.AutoModel.from_pretrained('aehrc/cxrmate-rrg24', trust_remote_code=True)
 transforms = v2.Compose(
     [
@@ -63,7 +64,7 @@ dataloader = DataLoader(dataset, batch_size=mbatch_size, shuffle=True)
 batch = next(iter(dataloader))
 
 output_ids = model.generate(
-    pixel_values=batch['images'],
+    pixel_values=batch['images'].to(device=device),
     max_length=512,
     num_beams=1,
     bad_words_ids=[[tokenizer.convert_tokens_to_ids('[NF]')], [tokenizer.convert_tokens_to_ids('[NI]')]],