henry commited on
Commit
0096d6d
·
1 Parent(s): ea39eba
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -25,7 +25,7 @@ def _load_model_tokenizer(checkpoint_path):
25
  torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
26
 
27
  model = AutoModelForCausalLM.from_pretrained(
28
- checkpoint_path, torch_dtype=torch_dtype, device_map=device_map, resume_download=True
29
  ).eval()
30
 
31
  model.generation_config.max_new_tokens = DEFAULT_MAX_NEW_TOKENS
 
25
  torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
26
 
27
  model = AutoModelForCausalLM.from_pretrained(
28
+ checkpoint_path, torch_dtype=torch_dtype, resume_download=True
29
  ).eval()
30
 
31
  model.generation_config.max_new_tokens = DEFAULT_MAX_NEW_TOKENS