Better gpu autoselection

hannah/train.py

@@ -85,8 +85,7 @@ def train(
             "gpu",
             "auto",
         ]:
-            if torch.cuda.is_available():
-                config.trainer.devices = auto_select_gpus(config.trainer.devices)
+            config.trainer.devices = auto_select_gpus(config.trainer.devices)
 
         if not config.trainer.fast_dev_run and not config.get("resume", False):
             clear_outputs()

hannah/utils/utils.py

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2023 Hannah contributors.
+# Copyright (c) 2024 Hannah contributors.
 #
 # This file is part of hannah.
 # See https://github.com/ekut-es/hannah for further info.
@@ -28,7 +28,7 @@
 import time
 from contextlib import _GeneratorContextManager, contextmanager
 from pathlib import Path
-from typing import Any, Callable, Iterator, List, Type, TypeVar
+from typing import Any, Callable, Iterator, List, Type, TypeVar, Union
 
 import hydra
 import numpy as np
@@ -213,7 +213,10 @@ def extract_from_download_cache(
         )
 
 
-def auto_select_gpus(gpus=1) -> List[int]:
+def auto_select_gpus(gpus=1) -> Union[List[int]]:
+    if not torch.cuda.is_available() or torch.cuda.device_count() < 1:
+        return gpus
+
     num_gpus = gpus
 
     if not nvsmi.is_nvidia_smi_on_path():