Spaces:

Chengxb888
/

test

Sleeping

Chengxb888 commited on Jul 18, 2024

Commit

0ed1ae4

verified ·

1 Parent(s): f4fa67b

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -17,7 +17,7 @@ def say_hello(msg: Annotated[str, Form()]):
     tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
     model = AutoModelForCausalLM.from_pretrained(
         "google/gemma-2b-it",
-        device_map="auto",
         torch_dtype=torch.bfloat16
     )
     print("token & msg")

     tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
     model = AutoModelForCausalLM.from_pretrained(
         "google/gemma-2b-it",
+        device_map="cpu",
         torch_dtype=torch.bfloat16
     )
     print("token & msg")