diff options
| author | AUTOMATIC <16777216c@gmail.com> | 2022-09-11 15:48:36 +0000 |
|---|---|---|
| committer | AUTOMATIC <16777216c@gmail.com> | 2022-09-11 15:48:36 +0000 |
| commit | f194457229e4537912467bc60ac3a873f473a63c (patch) | |
| tree | 93e9c445a79bb0f2e88e4058734594f318f1151c /modules/devices.py | |
| parent | 13008bab90305144591e0b2e233214a8b3415bba (diff) | |
| download | stable-diffusion-webui-gfx803-f194457229e4537912467bc60ac3a873f473a63c.tar.gz stable-diffusion-webui-gfx803-f194457229e4537912467bc60ac3a873f473a63c.tar.bz2 stable-diffusion-webui-gfx803-f194457229e4537912467bc60ac3a873f473a63c.zip | |
CLIP interrogator
Diffstat (limited to 'modules/devices.py')
| -rw-r--r-- | modules/devices.py | 16 |
1 files changed, 10 insertions, 6 deletions
diff --git a/modules/devices.py b/modules/devices.py index 25008a04..30d30b99 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -1,12 +1,16 @@ import torch - # has_mps is only available in nightly pytorch (for now), `getattr` for compatibility has_mps = getattr(torch, 'has_mps', False) +cpu = torch.device("cpu") + + def get_optimal_device(): - if torch.cuda.is_available(): - return torch.device("cuda") - if has_mps: - return torch.device("mps") - return torch.device("cpu") + if torch.cuda.is_available(): + return torch.device("cuda") + + if has_mps: + return torch.device("mps") + + return cpu |