chenjoya commited on
Commit
3188902
·
1 Parent(s): 58b8183
Files changed (1) hide show
  1. demo/infer.py +1 -1
demo/infer.py CHANGED
@@ -34,9 +34,9 @@ class LiveCCDemoInfer:
34
  def __init__(self, model_path: str = None, device_id: int = 0):
35
  self.model = Qwen2VLForConditionalGeneration.from_pretrained(
36
  model_path, torch_dtype="auto",
37
- device_map=f'cuda:{device_id}',
38
  attn_implementation='sdpa'
39
  )
 
40
  self.processor = AutoProcessor.from_pretrained(model_path, use_fast=False)
41
  self.streaming_eos_token_id = self.processor.tokenizer(' ...').input_ids[-1]
42
  self.model.prepare_inputs_for_generation = functools.partial(prepare_multiturn_multimodal_inputs_for_generation, self.model)
 
34
  def __init__(self, model_path: str = None, device_id: int = 0):
35
  self.model = Qwen2VLForConditionalGeneration.from_pretrained(
36
  model_path, torch_dtype="auto",
 
37
  attn_implementation='sdpa'
38
  )
39
+ self.model.to(f'cuda:{device_id}')
40
  self.processor = AutoProcessor.from_pretrained(model_path, use_fast=False)
41
  self.streaming_eos_token_id = self.processor.tokenizer(' ...').input_ids[-1]
42
  self.model.prepare_inputs_for_generation = functools.partial(prepare_multiturn_multimodal_inputs_for_generation, self.model)