fix: Manual selection of device

src/utils/utils.py

@@ -1,8 +1,8 @@
 from os import path
 
+import matplotlib.pyplot as plt
 import torch
 from torch.utils.data import TensorDataset
-import matplotlib.pyplot as plt
 
 
 def make_context_pairs(data: bytes, context_length: int) -> TensorDataset:
@@ -12,10 +12,12 @@ def make_context_pairs(data: bytes, context_length: int) -> TensorDataset:
     y = data[context_length:]
     return TensorDataset(x, y)
 
+
 def print_distribution(from_to: tuple[int, int], probabilities: list[float]):
     plt.hist(range(from_to[0], from_to[1]), weights=probabilities)
     plt.show()
 
+
 def print_losses(train_losses: list[float], validation_losses: list[float], filename: str | None = None, show=False):
     plt.plot(train_losses, label="Training loss")
     plt.plot(validation_losses, label="Validation loss")
@@ -33,6 +35,20 @@ def print_losses(train_losses: list[float], validation_losses: list[float], file
     plt.savefig(filename)
 
 
+def determine_device():
+    # NVIDIA GPUs (most HPC clusters)
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    # Apple Silicon (macOS)
+    elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+        return torch.device("mps")
+    # Intel GPUs (oneAPI)
+    elif hasattr(torch, "xpu") and torch.xpu.is_available():
+        return torch.device("xpu")
+    else:
+        return torch.device("cpu")
+
+
 def load_data(path: str) -> bytes:
     with open(path, "rb") as f:
         return f.read()