ZeroGPU-LLM-Inference

Running

Luigi commited on May 12

Commit

cc12a51

1 Parent(s): cb16ac6

revert device_map to audo

Files changed (1) hide show

app.py CHANGED Viewed

@@ -78,7 +78,7 @@ def load_pipeline(model_name):
                 tokenizer=tokenizer,
                 trust_remote_code=True,
                 torch_dtype=dtype,
-                device_map="xpu"
             )
             PIPELINES[model_name] = pipe
             return pipe
@@ -90,7 +90,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