Spaces:

aijadugar
/

ft_slm

Sleeping

aijadugar commited on 27 days ago

Commit

d706b30

verified ·

1 Parent(s): b172ee5

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -9,8 +9,9 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
     torch_dtype=torch.float16,
-    device_map="auto"
 )
 def generate_response(user_input):
     inputs = tokenizer(user_input, return_tensors="pt").to(model.device)

 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
     torch_dtype=torch.float16,
+    device_map=None
 )
+model = model.to("cpu")
 def generate_response(user_input):
     inputs = tokenizer(user_input, return_tensors="pt").to(model.device)