1

app.py

@@ -16,7 +16,7 @@ def predict(image):
     logits = outputs.logits
 
     # 결과를 이미지로 변환 (예: 가장 높은 확률을 가진 클래스 선택)
-    result = torch.argmax(logits, dim=1)
+    result = torch.argmax(logits)
     result = result.squeeze().cpu().numpy()
 
     # 여기에서는 단순화를 위해 결과를 그대로 반환합니다.