Update app.py

app.py

@@ -3,13 +3,13 @@ import re
 import gradio as gr
 from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel 
 
-
+device='cpu'
 encoder_checkpoint = "google/vit-base-patch16-224-in21k"
 decoder_checkpoint = "gpt2"
 model_checkpoint = "gagan3012/ViTGPT2I2A"
 feature_extractor = ViTFeatureExtractor.from_pretrained(encoder_checkpoint)
 tokenizer = AutoTokenizer.from_pretrained(decoder_checkpoint)
-model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint).to('cpu')
+model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint).to(device)
 
 
 def predict(image,max_length=64, num_beams=4):