Spaces:

TharanJ
/

skin-disease-classifier

Sleeping

TharanJ commited on Apr 19

Commit

a548af3

verified ·

1 Parent(s): cd4a6a0

edited return case

Files changed (1) hide show

model.py CHANGED Viewed

@@ -31,7 +31,7 @@ def predict(image):
     with torch.no_grad():
         outputs = model(**inputs)
         predicted_class = torch.argmax(outputs.logits, dim=1).item()
-    return predicted_class, class_names[predicted_class]
 # Create Gradio interface for user input

     with torch.no_grad():
         outputs = model(**inputs)
         predicted_class = torch.argmax(outputs.logits, dim=1).item()
+    return class_names[predicted_class]
 # Create Gradio interface for user input