Add get_device

src/utils/__init__.py

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

src/utils/get_device.py

@@ -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")