Spaces:

hshetty
/

self-driving-car

Runtime error

hshetty commited on Sep 27, 2022

Commit

afe1a22

1 Parent(s): da518b0

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -6,7 +6,7 @@ import torch
 extractor = AutoFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
 model = SegformerForSemanticSegmentation.from_pretrained("hshetty/my-segmentation-model")
 def classify(im):
-  inputs = extractor(images=im, return_tensors="pt").to("cuda")
   outputs = model(**inputs)
   logits = outputs.logits
   classes = logits[0].detach().cpu().numpy().argmax(axis=0)

 extractor = AutoFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
 model = SegformerForSemanticSegmentation.from_pretrained("hshetty/my-segmentation-model")
 def classify(im):
+  inputs = extractor(images=im, return_tensors="pt")
   outputs = model(**inputs)
   logits = outputs.logits
   classes = logits[0].detach().cpu().numpy().argmax(axis=0)