yachay commited on
Commit
8dff62f
·
1 Parent(s): 626813f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
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