Spaces:

Omnibus-archive
/

starcoderbase

Sleeping

Omnibus commited on Sep 7, 2023

Commit

a352172

1 Parent(s): b7c0853

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -6,6 +6,8 @@ import torch
 def fmt_prompt(prompt: str) -> str:
     return f"""[Instructions]:\n{prompt}\n\n[Response]:"""
 model_name = "abacaj/starcoderbase-1b-sft"
@@ -15,7 +17,7 @@ model = (
     transformers.AutoModelForCausalLM.from_pretrained(
         model_name,
     )
-    .to("cuda:0")
     .eval()
 )
 def chat_fn(prompt):

 def fmt_prompt(prompt: str) -> str:
     return f"""[Instructions]:\n{prompt}\n\n[Response]:"""
+#device = "cuda:0"
+device = "cpu"
 model_name = "abacaj/starcoderbase-1b-sft"
     transformers.AutoModelForCausalLM.from_pretrained(
         model_name,
     )
+    .to(device)
     .eval()
 )
 def chat_fn(prompt):