Updated from colab

app.py

@@ -8,10 +8,10 @@ extractor = AutoFeatureExtractor.from_pretrained("brendenc/my-segmentation-model
 model = AutoModelForImageClassification.from_pretrained("brendenc/my-segmentation-model")
 
 def classify(im):
-  inputs = extractor(images=im, return_tensors="pt").to("cuda")
+  inputs = extractor(images=im, return_tensors="pt")
   outputs = model(**inputs)
   logits = outputs.logits
-  classes = logits[0].detach().cpu().numpy().argmax(axis=0)
+  classes = logits[0].detach().numpy().argmax(axis=0)
   colors = np.array([[128,0,0], [128,128,0], [0, 0, 128], 	[128,0,128], [0, 0, 0]])
   return colors[classes]