ibraheemmoosa commited on
Commit
891b9bc
1 Parent(s): df59c87

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -35,7 +35,7 @@ model = MTRanker.from_pretrained('ibraheemmoosa/mt-ranker-base')
35
  def predict(source, translation1, translation2):
36
  model_input = "Source: {} Translation 0: {} Translation 1: {}".format(source, translation1, translation2)
37
  inputs = tokenizer([model_input], max_length=512, padding='max_length', truncation=True, return_tensors='pt')
38
- with autocast(dtype=torch.bfloat16):
39
  logits = model(inputs.input_ids, inputs.attention_mask)
40
  output_scores = torch.softmax(logits, dim=1)
41
  output_scores = output_scores[0]
 
35
  def predict(source, translation1, translation2):
36
  model_input = "Source: {} Translation 0: {} Translation 1: {}".format(source, translation1, translation2)
37
  inputs = tokenizer([model_input], max_length=512, padding='max_length', truncation=True, return_tensors='pt')
38
+ with torch.inference_mode():
39
  logits = model(inputs.input_ids, inputs.attention_mask)
40
  output_scores = torch.softmax(logits, dim=1)
41
  output_scores = output_scores[0]