andras/bipolaroid
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add get_device

Browse source
This commit is contained in:
Andras Schmelczer 2024-06-30 16:15:34 +01:00
commit 10249e6d26
No known key found for this signature in database
GPG key ID: FC8F2C3D3D1A718C
2 changed files with 8 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

1
src/utils/__init__.py
View file

@ -4,3 +4,4 @@ from .generate_rotation_matrices import generate_rotation_matrices
from .kldiv import kldiv from .kldiv import kldiv
from .set_up_logging import set_up_logging from .set_up_logging import set_up_logging
from .serialise_hparams import serialise_hparams from .serialise_hparams import serialise_hparams
from .get_device import get_device

7
src/utils/get_device.py Normal file
View file

@ -0,0 +1,7 @@
import torch
import os
def get_device() -> torch.device:
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
return torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")