Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

enable Pipeline to get device from model #30534

Merged
merged 12 commits into from
May 13, 2024
Merged

enable Pipeline to get device from model #30534

Changes from 3 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
ae6b3be
check model.device
faaany Apr 29, 2024
b3113ba
fix
faaany Apr 29, 2024
03a87fb
style fix
faaany Apr 29, 2024
be79f66
move model device
faaany Apr 30, 2024
dea4dbb
remove print
faaany Apr 30, 2024
385e8bd
add comment
faaany Apr 30, 2024
c1f9440
fix
faaany Apr 30, 2024
5722822
Merge branch 'huggingface:main' into device
faaany May 10, 2024
8347472
add unit test
faaany May 11, 2024
868c1ee
optimize
faaany May 11, 2024
f80667d
change test names and add more cases
faaany May 11, 2024
ff340cb
Update tests/pipelines/test_pipelines_common.py
faaany May 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/transformers/pipelines/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -830,6 +830,7 @@ def __init__(

# `accelerate` device map
hf_device_map = getattr(self.model, "hf_device_map", None)
model_device = getattr(self.model, "device", None)

if hf_device_map is not None and device is not None:
raise ValueError(
Expand All @@ -841,6 +842,8 @@ def __init__(
if hf_device_map is not None:
# Take the first device used by `accelerate`.
device = next(iter(hf_device_map.values()))
elif model_device is not None:
device = model_device
else:
device = -1

Expand Down