fix for tf-directml

iperov · May 4, 2022 · 0d19d8e · 0d19d8e
1 parent b9e571d
commit 0d19d8e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/core/leras/device.py b/core/leras/device.py
@@ -161,6 +161,7 @@ def initialize_main_env():
         if 'CUDA_VISIBLE_DEVICES' in os.environ.keys():
             os.environ.pop('CUDA_VISIBLE_DEVICES')
 
+        os.environ['TF_DIRECTML_KERNEL_CACHE_SIZE'] = '2500'
         os.environ['CUDA_CACHE_MAXSIZE'] = '2147483647'
         os.environ['TF_MIN_GPU_MULTIPROCESSOR_COUNT'] = '2'
         os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # tf log errors only