Update app.py

app.py

@@ -9,15 +9,15 @@ Original file is located at
 
 from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer
 import gradio as gr
-#import torch
+import torch
 from PIL import Image
 
 model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 feature_extractor = ViTFeatureExtractor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 
-# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# model.to(device)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)