Support all NVidia devices (#143)

metal-stack · Oct 15, 2024 · eaceba8 · eaceba8
1 parent 9a01121
commit eaceba8
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/cmd/register/register.go b/cmd/register/register.go
@@ -245,12 +245,12 @@ func (r *Register) detectGPUs() (pci.Devices, error) {
 		if !strings.Contains(strings.ToLower(device.VendorName), "nvidia") {
 			continue
 		}
+		r.log.Info("add gpu", "vendor", device.VendorName, "device", device.DeviceName, "classname", device.ClassName)
 
-		// TODO if new models must be supported, this code must be refactored
-		if strings.Contains(strings.ToLower(device.DeviceName), "rtx") {
-			r.log.Info("add gpu", "vendor", device.VendorName, "device", device.DeviceName)
-			result = append(result, device)
+		if !strings.Contains(strings.ToLower(device.ClassName), "display") {
+			continue
 		}
+		result = append(result, device)
 	}
 
 	return result, nil