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

feat: allow cpu device even if gpu available #3417

Merged
merged 2 commits into from
Mar 29, 2024
Merged

feat: allow cpu device even if gpu available #3417

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f074b98
feat: allow cpu device even if gpu available
drbh Mar 8, 2024
cc1d938
fix: prefer `device_id` name
drbh Mar 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 6 additions & 5 deletions flair/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,15 +18,16 @@
device: torch.device
"""Flair is using a single device for everything. You can set this device by overwriting this variable.

This value will be automatically set to the first found GPU if available and to CPU otherwise.
You can choose a specific GPU, by setting the `FLAIR_DEVICE` environment variable to its index.
The device will be automatically set to the first available GPU if a GPU is present and the 'FLAIR_DEVICE' environment
variable is not set to 'cpu', otherwise it will default to the CPU, and a specific GPU can be chosen by setting the 'FLAIR_DEVICE'
environment variable to its index.
"""

# Get the device from the environment variable
device_id = os.environ.get("FLAIR_DEVICE")

# global variable: device
if torch.cuda.is_available():
device_id = os.environ.get("FLAIR_DEVICE")

if torch.cuda.is_available() and device_id != "cpu":
# No need for correctness checks, torch is doing it
device = torch.device(f"cuda:{device_id}") if device_id else torch.device("cuda:0")
else:
Expand Down
Loading