Spaces:
Running
Running
edited return case
Browse files
model.py
CHANGED
@@ -31,7 +31,7 @@ def predict(image):
|
|
31 |
with torch.no_grad():
|
32 |
outputs = model(**inputs)
|
33 |
predicted_class = torch.argmax(outputs.logits, dim=1).item()
|
34 |
-
return
|
35 |
|
36 |
|
37 |
# Create Gradio interface for user input
|
|
|
31 |
with torch.no_grad():
|
32 |
outputs = model(**inputs)
|
33 |
predicted_class = torch.argmax(outputs.logits, dim=1).item()
|
34 |
+
return class_names[predicted_class]
|
35 |
|
36 |
|
37 |
# Create Gradio interface for user input
|