salaz055 commited on
Commit
afa3b98
1 Parent(s): c94d411

Update app.py

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