diff --git a/sdk_container/src/third_party/coreos-overlay/x11-drivers/nvidia-drivers/files/bin/setup-nvidia b/sdk_container/src/third_party/coreos-overlay/x11-drivers/nvidia-drivers/files/bin/setup-nvidia
index 0c6f52e158..6ecb16d334 100644
--- a/sdk_container/src/third_party/coreos-overlay/x11-drivers/nvidia-drivers/files/bin/setup-nvidia
+++ b/sdk_container/src/third_party/coreos-overlay/x11-drivers/nvidia-drivers/files/bin/setup-nvidia
@@ -136,10 +136,15 @@ function verify_installation() {
 }
 
 function is_nvidia_probe_required() {
-  if [[ -z "$(lspci | grep -i "${NVIDIA_PRODUCT_TYPE}")" ]]; then
-    return 1
+  # Vendor: NVIDIA, Class: VGA compatible controller
+  if [[ -n "$(lspci -d 10de:*:0300)" ]]; then
+    return 0
   fi
-  return 0
+  # Vendor: NVIDIA, Class: 3D controller
+  if [[ -n "$(lspci -d 10de:*:0302)" ]]; then
+    return 0
+  fi
+  return 1
 }
 
 function is_nvidia_installation_required() {
diff --git a/sdk_container/src/third_party/coreos-overlay/x11-drivers/nvidia-drivers/nvidia-drivers-535.104.05.ebuild b/sdk_container/src/third_party/coreos-overlay/x11-drivers/nvidia-drivers/nvidia-drivers-535.104.05-r1.ebuild
similarity index 100%
rename from sdk_container/src/third_party/coreos-overlay/x11-drivers/nvidia-drivers/nvidia-drivers-535.104.05.ebuild
rename to sdk_container/src/third_party/coreos-overlay/x11-drivers/nvidia-drivers/nvidia-drivers-535.104.05-r1.ebuild