LisaMegaWatts commited on
Commit
fa396cb
·
verified ·
1 Parent(s): 11261f7

fix: use dtype= instead of deprecated torch_dtype=

Browse files
Files changed (1) hide show
  1. server.py +1 -1
server.py CHANGED
@@ -151,7 +151,7 @@ def generate_streaming(prompt: str, max_tokens: int = 200,
151
  print(f"Loading base model: {BASE_MODEL_REPO} ...")
152
  BASE_MODEL = AutoModelForCausalLM.from_pretrained(
153
  BASE_MODEL_REPO,
154
- torch_dtype=torch.float32,
155
  device_map="cpu",
156
  low_cpu_mem_usage=True,
157
  )
 
151
  print(f"Loading base model: {BASE_MODEL_REPO} ...")
152
  BASE_MODEL = AutoModelForCausalLM.from_pretrained(
153
  BASE_MODEL_REPO,
154
+ dtype=torch.float32,
155
  device_map="cpu",
156
  low_cpu_mem_usage=True,
157
  )