Update README.md

README.md

@@ -11,7 +11,8 @@ Here is how to use this model to caption an image of the Flickr8k dataset:
 import torch
 import requests
 from PIL import Image
-from transformers import ViTFeatureExtractor, AutoTokenizer, VisionEncoderDecoderModel
+from transformers import ViTFeatureExtractor, AutoTokenizer, 
+                         VisionEncoderDecoderModel
 
 if torch.cuda.is_available():
     device = 'cuda'
@@ -23,10 +24,10 @@ image = Image.open(requests.get(url, stream=True).raw)
 
 encoder_checkpoint = 'google/vit-base-patch16-224'
 decoder_checkpoint = 'surajp/gpt2-hindi'
-
+model_checkpoint = 'team-indain-image-caption/hindi-image-captioning'
 feature_extractor = ViTFeatureExtractor.from_pretrained(encoder_checkpoint)
 tokenizer = AutoTokenizer.from_pretrained(decoder_checkpoint)
-model = VisionEncoderDecoderModel.from_pretrained('team-indain-image-caption/hindi-image-captioning').to(device)
+model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint).to(device)
 
 #Inference
 sample = feature_extractor(image, return_tensors="pt").pixel_values.to(device)