SRDdev commited on
Commit
191c3b2
·
1 Parent(s): d3f1dad

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -5,10 +5,10 @@ from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecode
5
 
6
 
7
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
8
- model.to(device)
9
  encoder_checkpoint = 'google/vit-base-patch16-224'
10
  decoder_checkpoint = 'sachin/vit2distilgpt2'
11
  model_checkpoint = 'nlpconnect/vit-gpt2-image-captioning'
 
12
  feature_extractor = ViTFeatureExtractor.from_pretrained(encoder_checkpoint)
13
  tokenizer = AutoTokenizer.from_pretrained(decoder_checkpoint)
14
  model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint).to(device)
 
5
 
6
 
7
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
8
  encoder_checkpoint = 'google/vit-base-patch16-224'
9
  decoder_checkpoint = 'sachin/vit2distilgpt2'
10
  model_checkpoint = 'nlpconnect/vit-gpt2-image-captioning'
11
+ model.to(device)
12
  feature_extractor = ViTFeatureExtractor.from_pretrained(encoder_checkpoint)
13
  tokenizer = AutoTokenizer.from_pretrained(decoder_checkpoint)
14
  model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint).to(device)