Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -68,7 +68,7 @@ class ByT5ForTextGeotagging(PreTrainedModel):
|
|
68 |
if return_coordinates:
|
69 |
class_idx = torch.argmax(logits, dim=1).item()
|
70 |
coordinates = self.config.class_to_location.get(str(class_idx))
|
71 |
-
confidence = torch.max(logits).item()
|
72 |
return logits, coordinates, confidence
|
73 |
else:
|
74 |
return logits
|
|
|
68 |
if return_coordinates:
|
69 |
class_idx = torch.argmax(logits, dim=1).item()
|
70 |
coordinates = self.config.class_to_location.get(str(class_idx))
|
71 |
+
confidence = torch.max(torch.nn.functional.softmax(logits)).item()
|
72 |
return logits, coordinates, confidence
|
73 |
else:
|
74 |
return logits
|