fix: Manual selection of device

main.py

@@ -1,17 +1,13 @@
-import torch
-
 from src.args import parse_arguments
 from src.process import compress
 from src.train import train
+from src.utils import determine_device
 
 
 def main():
     args, print_help = parse_arguments()
 
-    if torch.accelerator.is_available():
-        device = torch.accelerator.current_accelerator().type
-    else:
-        device = "cpu"
+    device = determine_device()
     print(f"Running on device: {device}...")
 
     match args.mode: