[fix] device_lib.list_local_devices() may run fail in some case, then…

… use CPU as default device.
tensorflow · Sep 17, 2024 · 07cdb8e · 07cdb8e
1 parent e85da20
commit 07cdb8e
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py
@@ -541,11 +541,14 @@ def default_partition_fn(keys, shard_num):
     self.short_file_name = short_file_name
 
     def _get_default_devices():
-      gpu_list = [
-          x.name
-          for x in device_lib.list_local_devices()
-          if x.device_type == "GPU"
-      ]
+      try:
+        gpu_list = [
+            x.name
+            for x in device_lib.list_local_devices()
+            if x.device_type == "GPU"
+        ]
+      except:
+        gpu_list = []
       return gpu_list[0:1] or [
           "/CPU:0",
       ]