hshetty commited on
Commit
afe1a22
1 Parent(s): da518b0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -6,7 +6,7 @@ import torch
6
  extractor = AutoFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
7
  model = SegformerForSemanticSegmentation.from_pretrained("hshetty/my-segmentation-model")
8
  def classify(im):
9
- inputs = extractor(images=im, return_tensors="pt").to("cuda")
10
  outputs = model(**inputs)
11
  logits = outputs.logits
12
  classes = logits[0].detach().cpu().numpy().argmax(axis=0)
 
6
  extractor = AutoFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
7
  model = SegformerForSemanticSegmentation.from_pretrained("hshetty/my-segmentation-model")
8
  def classify(im):
9
+ inputs = extractor(images=im, return_tensors="pt")
10
  outputs = model(**inputs)
11
  logits = outputs.logits
12
  classes = logits[0].detach().cpu().numpy().argmax(axis=0)