Fix GPU flag for A100 (#162)

dask · Nov 11, 2020 · 8002e0b · 8002e0b
1 parent 09c2897
commit 8002e0b
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py
@@ -168,12 +168,16 @@ def create_gcp_config(self):
 
         return config
 
+    @property
+    def gpu_instance(self):
+        return "gpu" in self.machine_type or bool(self.ngpus)
+
     async def create_vm(self):
 
         self.cloud_init = self.cluster.render_cloud_init(
             image=self.docker_image,
             command=self.command,
-            gpu_instance=bool(self.ngpus),
+            gpu_instance=self.gpu_instance,
             bootstrap=self.bootstrap,
             auto_shutdown=self.cluster.auto_shutdown,
             env_vars=self.env_vars,