fix device request logging

a-r-j · Nov 13, 2023 · 5315f7e · 5315f7e
1 parent fc7521c
commit 5315f7e
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/proteinworkshop/configs/config.py b/proteinworkshop/configs/config.py
@@ -245,18 +245,20 @@ def validate_cuda(cfg: DictConfig) -> DictConfig:
     # if cfg.trainer.devices <= 1:
     # cfg.trainer.strategy = None
 
-    logger.debug(f"Requested GPUs: {cfg.get('trainer.devices')}.")
-    if isinstance(cfg.get("trainer.devices"), int):
+    logger.debug(f"Requested GPUs: {cfg.trainer.get('devices')}.")
+    if isinstance(cfg.trainer.get("devices"), int):
         cfg.trainer.devices = min(
             torch.cuda.device_count(), cfg.trainer.devices
         )
         logger.debug(f"GPU count set to: {cfg.trainer.devices}")
 
     requesting_multiple_device_indices = (
-        isinstance(cfg.get("trainer.devices"), list) and len(cfg.get("trainer.devices")) > 1
+        isinstance(cfg.trainer.get("devices"), list)
+        and len(cfg.trainer.get("devices")) > 1
     )
     requesting_multiple_devices = (
-        isinstance(cfg.get("trainer.devices"), int) and cfg.get("trainer.devices") > 1
+        isinstance(cfg.trainer.get("devices"), int)
+        and cfg.trainer.get("devices") > 1
     )
     if (
         requesting_multiple_device_indices or requesting_multiple_devices