chenjoya commited on
Commit
fddefb0
·
verified ·
1 Parent(s): dc317e6

Update demo/infer.py

Browse files
Files changed (1) hide show
  1. demo/infer.py +1 -1
demo/infer.py CHANGED
@@ -35,7 +35,7 @@ class LiveCCDemoInfer:
35
  def __init__(self, model_path: str = None, device: str = 'cpu'):
36
  self.model = Qwen2VLForConditionalGeneration.from_pretrained(
37
  model_path, torch_dtype="auto",
38
- # device_map=f'cuda:{device_id}',
39
  # attn_implementation='flash_attention_2'
40
  )
41
  self.model.to(device)
 
35
  def __init__(self, model_path: str = None, device: str = 'cpu'):
36
  self.model = Qwen2VLForConditionalGeneration.from_pretrained(
37
  model_path, torch_dtype="auto",
38
+ device_map=device,
39
  # attn_implementation='flash_attention_2'
40
  )
41
  self.model.to(device)