Fix "Not creating XLA devices, tf_xla_enable_xla_devices not set"

lightvector · May 3, 2022 · b2d4ae9 · b2d4ae9
1 parent 895f81d
commit b2d4ae9
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/train.py b/python/train.py
@@ -152,6 +152,9 @@ def trainlog(s):
     multi_gpu_device_ids.append("/GPU:" + str(int(piece)))
   num_gpus_used = len(multi_gpu_device_ids)
 
+# Fix for tensorflow 2.4: Not creating XLA devices, tf_xla_enable_xla_devices not set
+os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
+
 
 # MODEL ----------------------------------------------------------------
 printed_model_yet = False