syedfaisalabrar commited on
Commit
826719d
·
verified ·
1 Parent(s): f520e17

Update frontPrompt.py

Browse files
Files changed (1) hide show
  1. frontPrompt.py +2 -2
frontPrompt.py CHANGED
@@ -32,8 +32,8 @@ def main(image_path):
32
  path = "OpenGVLab/InternVL2_5-4B"
33
  model = AutoModel.from_pretrained(
34
  path,
35
- torch_dtype=torch.bfloat16,
36
- load_in_8bit=True,
37
  low_cpu_mem_usage=True,
38
  use_flash_attn=True,
39
  trust_remote_code=True).eval()
 
32
  path = "OpenGVLab/InternVL2_5-4B"
33
  model = AutoModel.from_pretrained(
34
  path,
35
+ torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
36
+ # load_in_8bit=True,
37
  low_cpu_mem_usage=True,
38
  use_flash_attn=True,
39
  trust_remote_code=True).eval()