guangyil commited on
Commit
0b8a6d2
·
verified ·
1 Parent(s): bc6a876

Update infer.py

Browse files
Files changed (1) hide show
  1. infer.py +2 -0
infer.py CHANGED
@@ -61,6 +61,8 @@ def is_audio_output_task(task_type):
61
 
62
  def eval_model(model, tokenizer, tokenizer_voila, model_type, task_type, history, ref_embs, ref_embs_mask, max_new_tokens=512):
63
  # step1: initializing
 
 
64
  num_codebooks = model.config.num_codebooks
65
  codebook_size = model.config.codebook_size
66
 
 
61
 
62
  def eval_model(model, tokenizer, tokenizer_voila, model_type, task_type, history, ref_embs, ref_embs_mask, max_new_tokens=512):
63
  # step1: initializing
64
+ model.to('cuda')
65
+ tokenizer_voila.to('cuda')
66
  num_codebooks = model.config.num_codebooks
67
  codebook_size = model.config.codebook_size
68