Automatically assign available hardware on the fly, in-line with PyTorch code.
Project description
autodevice
Automatically assign devices in-line with pytorch code
Usage
from autodevice import AutoDevice
x = torch.randn([200, 50]).to(AutoDevice())
CUDA/GPU:
tensor([[ 2.6905, -0.3037, -0.3607],
[ 0.2258, -0.1755, 0.6599],
[ 1.3046, -0.9389, 0.7358]], device='cuda:0')
CPU:
tensor([[ 2.6905, -0.3037, -0.3607],
[ 0.2258, -0.1755, 0.6599],
[ 1.3046, -0.9389, 0.7358]])
On Apple Silicon (M1, M2):
tensor([[ 0.5382, 1.1173, 1.1175],
[-0.0125, -0.2406, 0.2343],
[-0.6067, -0.7728, 0.1697]], device='mps:0')
Installation
pip install autodevice
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
autodevice-0.1.0rc0.tar.gz
(3.8 kB
view hashes)
Built Distribution
Close
Hashes for autodevice-0.1.0rc0-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 0832da9b14a9b7fac0cd976848df780eaa57f799f3a90fd4f682c183a0dfc90d |
|
MD5 | 69fa6b374ea387f189259c5afaea7aac |
|
BLAKE2b-256 | b7abd602de31001328169b7a35b750bb32c9d3b90d1eefbe3ae838b0c7353afb |