wehub-resource-sync
cddb07a176
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
21 lines
772 B
Python
21 lines
772 B
Python
import itertools
|
|
|
|
import torch
|
|
|
|
|
|
def get_effective_device(model: torch.nn.Module) -> torch.device:
|
|
"""A utility to infer the 'effective' device of a model.
|
|
|
|
This utility handles the case where a model is partially loaded onto the GPU, so is safer than just calling:
|
|
`next(iter(model.parameters())).device`.
|
|
|
|
In the worst case, this utility has to check all model parameters, so if you already know the intended model device,
|
|
then it is better to avoid calling this function.
|
|
"""
|
|
# If all parameters are on the CPU, return the CPU device. Otherwise, return the first non-CPU device.
|
|
for p in itertools.chain(model.parameters(), model.buffers()):
|
|
if p.device.type != "cpu":
|
|
return p.device
|
|
|
|
return torch.device("cpu")
|