tianyang commited on
Commit
92828cc
·
1 Parent(s): 3a5a33e

change fp16 to bf16 since no GPU

Browse files
Files changed (1) hide show
  1. utils/inference.py +2 -2
utils/inference.py CHANGED
@@ -48,14 +48,14 @@ def load_tokenizer_and_model(base_model, adapter_model, load_8bit=True):
48
  base_model,
49
  device_map={"": device},
50
  low_cpu_mem_usage=True,
51
- torch_dtype=torch.float16,
52
  offload_folder="."
53
  )
54
  if adapter_model is not None:
55
  model = PeftModel.from_pretrained(
56
  model,
57
  adapter_model,
58
- torch_dtype=torch.float16,
59
  offload_folder="."
60
  )
61
 
 
48
  base_model,
49
  device_map={"": device},
50
  low_cpu_mem_usage=True,
51
+ torch_dtype=torch.bfloat16,
52
  offload_folder="."
53
  )
54
  if adapter_model is not None:
55
  model = PeftModel.from_pretrained(
56
  model,
57
  adapter_model,
58
+ torch_dtype=torch.bfloat16,
59
  offload_folder="."
60
  )
61