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.0.2.tar.gz
(3.0 kB
view hashes)
Built Distribution
Close
Hashes for autodevice-0.0.2-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | d665ab0f87bb3f2acddd57547f114ac216e8322b7ac9640c4c33edd32a6eb44d |
|
MD5 | fbdf1a5a4e48b1870b353ec8743b1911 |
|
BLAKE2b-256 | 15b551546eeb55750b5b741649094b1fcb3ce3419a489d6e00250b649fe6d9ab |