kanhatakeyama commited on
Commit
69bf2b0
1 Parent(s): b6dcc1d

Upload model

Browse files
Files changed (1) hide show
  1. MoEModel.py +4 -1
MoEModel.py CHANGED
@@ -73,8 +73,11 @@ class MoeModel(PreTrainedModel):
73
  return ret
74
 
75
 
 
 
 
76
  def perplexity(model, tokenized_input) -> torch.Tensor:
77
  with torch.inference_mode():
78
- output = model(tokenized_input, labels=tokenized_input)
79
  ppl = torch.exp(output.loss)
80
  return ppl.item()
 
73
  return ret
74
 
75
 
76
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
77
+
78
+
79
  def perplexity(model, tokenized_input) -> torch.Tensor:
80
  with torch.inference_mode():
81
+ output = model(tokenized_input.to(device), labels=tokenized_input)
82
  ppl = torch.exp(output.loss)
83
  return ppl.item()