-
Notifications
You must be signed in to change notification settings - Fork 37
Define XLA.DEFAULT_DEVICE to set default device ID
#1866
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.
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
base: main
Are you sure you want to change the base?
Conversation
386197e to
2c169ef
Compare
444cf6b to
3f9be82
Compare
|
To demonstrate this works, on top of #1860 I get julia> using Reactant
julia> Reactant.XLA.default_device()
Reactant.XLA.PJRT.Device(Ptr{Nothing}(0x0000000042369bd0), "TT:0 Wormhole_b0")
julia> ENV["REACTANT_DEFAULT_DEVICE"] = "31"
"31"
julia> Reactant.XLA.default_device()
Reactant.XLA.PJRT.Device(Ptr{Nothing}(0x000000004d618490), "TT:31 Wormhole_b0") |
REACTANT_DEFAULT_DEVICE to set default device ID with env varXLA.DEFAULT_DEVICE to set default device ID
| 0-based index of default device to use, by default 0 (first available device). | ||
| """ | ||
| const DEFAULT_DEVICE = Ref{Int}(0) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
if we're not doing an env var we can leave this as 1-indexed. [reason for env var is cuz cuda_visible_devices and friends are 0-indexed]
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
also if you want we can still do DEFAULT_DEVICE = Ref{Int}(something(tryparse, ...)
No description provided.