From 0d19d8ec8ed2def5954dd8061ded64551efe270a Mon Sep 17 00:00:00 2001 From: iperov Date: Wed, 4 May 2022 16:41:17 +0400 Subject: [PATCH] fix for tf-directml --- core/leras/device.py | 1 + 1 file changed, 1 insertion(+) diff --git a/core/leras/device.py b/core/leras/device.py index 31d2f88..a2ba371 100644 --- a/core/leras/device.py +++ b/core/leras/device.py @@ -161,6 +161,7 @@ class Devices(object): 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