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.0.tar.gz
(3.8 kB
view hashes)
Built Distribution
Close
Hashes for autodevice-0.1.0-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 5c5f57c45b40cd4af62978abd3f5beef5b2f088f506f5c729c3bce5d70052eca |
|
MD5 | e47c84bd6cf9a5d3c54fa5b3f97095d7 |
|
BLAKE2b-256 | cd79fc07ca3758b683e42417a47be229b72110d10f5e0e7985458e02c1d88fc8 |