ZeroGPU-LLM-Inference

Running

Luigi commited on May 2

Commit

1242438

1 Parent(s): 2882063

bugfix: set device to xpu by mistake

Files changed (1) hide show

app.py CHANGED Viewed

@@ -70,7 +70,7 @@ def load_pipeline(model_name):
                 tokenizer=tokenizer,
                 trust_remote_code=True,
                 torch_dtype=dtype,
-                device_map="xpu"
             )
             PIPELINES[model_name] = pipe
             return pipe
@@ -82,7 +82,7 @@ def load_pipeline(model_name):
         model=repo,
         tokenizer=tokenizer,
         trust_remote_code=True,
-        device_map="xpu"
     )
     PIPELINES[model_name] = pipe
     return pipe

                 tokenizer=tokenizer,
                 trust_remote_code=True,
                 torch_dtype=dtype,
+                device_map="auto"
             )
             PIPELINES[model_name] = pipe
             return pipe
         model=repo,
         tokenizer=tokenizer,
         trust_remote_code=True,
+        device_map="auto"
     )
     PIPELINES[model_name] = pipe
     return pipe