Add compute capability resolver

Signed-off-by: Jan Baczek <[email protected]>

nemo/core/config/hydra_runner.py

@@ -23,6 +23,21 @@
 from hydra.types import TaskFunction
 from omegaconf import DictConfig, OmegaConf
 
+def _get_gpu_name():
+    import pynvml
+    pynvml.nvmlInit()
+    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+    cuda_capability, _ = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
+    pynvml.nvmlShutdown()
+    if cuda_capability == 8:
+        return "a100"
+    elif cuda_capability == 9:
+        return "h100"
+    else:
+        return None
+
+omegaconf.OmegaConf.register_new_resolver("gpu_name", _get_gpu_name)
+
 # multiple interpolated values in the config
 OmegaConf.register_new_resolver("multiply", lambda x, y: x * y)