shmuelamar commited on
Commit
61c4504
·
unverified ·
1 Parent(s): 9126461

fix use after free

Browse files
Files changed (1) hide show
  1. app.py +4 -1
app.py CHANGED
@@ -102,12 +102,15 @@ def completion(prompt: str, model_id: str):
102
  temperature=None,
103
  )
104
 
 
 
 
105
  # cleanup memory
106
  del model, tokenizer
107
  torch.cuda.empty_cache()
108
  gc.collect()
109
 
110
- return tokenizer.decode(outputs[0][input_ids.shape[-1] :], skip_special_tokens=True)
111
 
112
 
113
  def completion_openrouter(prompt: str, model_id: str):
 
102
  temperature=None,
103
  )
104
 
105
+ # decode response
106
+ resp = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
107
+
108
  # cleanup memory
109
  del model, tokenizer
110
  torch.cuda.empty_cache()
111
  gc.collect()
112
 
113
+ return resp
114
 
115
 
116
  def completion_openrouter(prompt: str, model_id: str):