Update README.md

README.md

@@ -1,7 +1,5 @@
 ---
 license: mit
-datasets:
-- ydshieh/coco_dataset_script
 language:
 - en
 metrics:
@@ -26,8 +24,9 @@ This is a sequence to sequence model for image-captioning. The architecture is <
 The way to download and run this: 
 ## Your code here
 ```python
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 from transformers import pipeline
-image_captioner = pipeline("image-to-text", model="sadassa17/rgb-language_cap", max_new_tokens=200)
+image_captioner = pipeline("image-to-text", model="sadassa17/rgb-language_cap", max_new_tokens=200, device=device)
 filename = 'path/to/file'
 generated_captions = image_captioner(filename)
 print(generated_captions)