diff --git a/test/e2e/framework/gpu_util.go b/test/e2e/framework/gpu_util.go
index d0ff9798f19..2cd9f33ee77 100644
--- a/test/e2e/framework/gpu_util.go
+++ b/test/e2e/framework/gpu_util.go
@@ -32,7 +32,7 @@ const (
 	// TODO: Parametrize it by making it a feature in TestFramework.
 	// so we can override the daemonset in other setups (non COS).
 	// GPUDevicePluginDSYAML is the official Google Device Plugin Daemonset NVIDIA GPU manifest for GKE
-	GPUDevicePluginDSYAML = "https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/device-plugin-daemonset.yaml"
+	GPUDevicePluginDSYAML = "https://raw.githubusercontent.com/kubernetes/kubernetes/master/cluster/addons/device-plugins/nvidia-gpu/daemonset.yaml"
 )
 
 // TODO make this generic and not linked to COS only
@@ -61,8 +61,8 @@ func NVIDIADevicePlugin(ns string) *v1.Pod {
 
 		Spec: ds.Spec.Template.Spec,
 	}
-	// Remove NVIDIA drivers installation
-	p.Spec.InitContainers = []v1.Container{}
+	// Remove node affinity
+	p.Spec.Affinity = nil
 
 	return p
 }
diff --git a/test/e2e/scheduling/nvidia-gpus.go b/test/e2e/scheduling/nvidia-gpus.go
index 5deeea25f7e..44055575f96 100644
--- a/test/e2e/scheduling/nvidia-gpus.go
+++ b/test/e2e/scheduling/nvidia-gpus.go
@@ -163,7 +163,7 @@ func testNvidiaGPUsOnCOS(f *framework.Framework) {
 	framework.Logf("Cluster is running on COS. Proceeding with test")
 
 	if f.BaseName == "device-plugin-gpus" {
-		dsYamlUrl = framework.GPUDevicePluginDSYAML
+		dsYamlUrl = "https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/daemonset.yaml"
 		gpuResourceName = framework.NVIDIAGPUResourceName
 		podCreationFunc = makeCudaAdditionDevicePluginTestPod
 	} else {