Caslow commited on
Commit
9efc447
·
1 Parent(s): 12bb3cd
Files changed (1) hide show
  1. inference.py +2 -1
inference.py CHANGED
@@ -29,9 +29,10 @@ def load_model(
29
 
30
  model = AutoModelForCausalLM.from_pretrained(
31
  pretrained_model_name_or_path=model_name,
 
32
  torch_dtype=torch.float32, # Use float32 for CPU
33
  low_cpu_mem_usage=True # Helps with memory efficiency
34
- ).to(device)
35
 
36
  model.eval() # Set model to evaluation mode
37
 
 
29
 
30
  model = AutoModelForCausalLM.from_pretrained(
31
  pretrained_model_name_or_path=model_name,
32
+ device_map="cpu",
33
  torch_dtype=torch.float32, # Use float32 for CPU
34
  low_cpu_mem_usage=True # Helps with memory efficiency
35
+ )
36
 
37
  model.eval() # Set model to evaluation mode
38