trminhnam20082002 commited on
Commit
171c344
1 Parent(s): 55e492d

feat: add model device before generating

Browse files
Files changed (1) hide show
  1. utils.py +2 -0
utils.py CHANGED
@@ -161,6 +161,8 @@ def make_input_sentence_from_strings(data):
161
  def generate_description(
162
  input_string, model, tokenizer, device, max_len, model_name, beam_size
163
  ):
 
 
164
  inputs = prepare_single_model_inputs(
165
  input_string, tokenizer, max_len=max_len, device=device
166
  )
 
161
  def generate_description(
162
  input_string, model, tokenizer, device, max_len, model_name, beam_size
163
  ):
164
+ model.eval()
165
+ model = model.to(device)
166
  inputs = prepare_single_model_inputs(
167
  input_string, tokenizer, max_len=max_len, device=device
168
  )