K00B404 commited on
Commit
6d26b7f
Β·
1 Parent(s): 7081866

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +7 -1
model.py CHANGED
@@ -18,7 +18,13 @@ if torch.cuda.is_available():
18
  use_safetensors=False,
19
  )
20
  else:
21
- model = None
 
 
 
 
 
 
22
  tokenizer = AutoTokenizer.from_pretrained(model_id)
23
 
24
 
 
18
  use_safetensors=False,
19
  )
20
  else:
21
+ # Switch PyTorch to CPU mode
22
+ torch.set_default_tensor_type(torch.FloatTensor)
23
+ model = AutoModelForCausalLM.from_pretrained(
24
+ model_id,
25
+ torch_dtype=torch.float32,
26
+ device_map='cpu')
27
+
28
  tokenizer = AutoTokenizer.from_pretrained(model_id)
29
 
30