aframson commited on
Commit
adba112
·
1 Parent(s): 8d90985
Files changed (1) hide show
  1. modelLM.py +1 -1
modelLM.py CHANGED
@@ -52,7 +52,7 @@ class OBILanguageModel(PreTrainedModel):
52
  else:
53
  loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), targets.view(-1))
54
 
55
- return logits, loss
56
 
57
 
58
  def generate(self, idx, max_new_tokens):
 
52
  else:
53
  loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), targets.view(-1))
54
 
55
+ return (logits, loss)
56
 
57
 
58
  def generate(self, idx, max_new_tokens):