fix: fallback to torch.cuda.mem_get_info when NVML memory query is unsupported

nvmlDeviceGetMemoryInfo returns NVML_ERROR_NOT_SUPPORTED on DGX Spark
(GB10). Fall back to torch.cuda.mem_get_info which works on all CUDA
devices.

Signed-off-by: Daniel Bustamante Ospina <dbustamante70@gmail.com>

Author: dbuos

nemo_rl/utils/nvml.py

@@ -78,13 +78,17 @@ def get_device_uuid(device_idx: int) -> str:
 
 
 def get_free_memory_bytes(device_idx: int) -> float:
-    """Get the free memory of a CUDA device in bytes using NVML."""
+    """Get the free memory of a CUDA device in bytes using NVML, with torch.cuda fallback."""
     global_device_idx = device_id_to_physical_device_id(device_idx)
     with nvml_context():
         try:
             handle = pynvml.nvmlDeviceGetHandleByIndex(global_device_idx)
             return pynvml.nvmlDeviceGetMemoryInfo(handle).free
-        except pynvml.NVMLError as e:
-            raise RuntimeError(
-                f"Failed to get free memory for device {device_idx} (global index: {global_device_idx}): {e}"
-            )
+        except pynvml.NVMLError:
+            pass
+
+    # Fallback for GPUs where NVML memory query is not supported (e.g. DGX Spark)
+    import torch
+
+    free, _total = torch.cuda.mem_get_info(device_idx)
+    return free