viethoangtranduong commited on
Commit
f2e31d0
1 Parent(s): 76c2f69

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +2 -1
handler.py CHANGED
@@ -153,6 +153,7 @@ class EndpointHandler():
153
  )
154
 
155
  output_strs = self.tokenizer.batch_decode(outputs.sequences[:, input_length:], skip_special_tokens=True)
156
- output_strs = [output_str.replace("\n<human>:", "") for output_str in output_strs]
157
 
 
158
  return {"generated_text": output_strs}
 
153
  )
154
 
155
  output_strs = self.tokenizer.batch_decode(outputs.sequences[:, input_length:], skip_special_tokens=True)
156
+ output_strs = [output_str.replace("\n<human>:", "") for output_str in output_strs]
157
 
158
+ torch.cuda.empty_cache()
159
  return {"generated_text": output_strs}