diff --git a/nnlib/nnlib.py b/nnlib/nnlib.py
index c8031bf..f2bd9ab 100644
--- a/nnlib/nnlib.py
+++ b/nnlib/nnlib.py
@@ -139,6 +139,8 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
 
         if 'CUDA_VISIBLE_DEVICES' in os.environ.keys():
             os.environ.pop('CUDA_VISIBLE_DEVICES')
+            
+        os.environ['CUDA_​CACHE_​MAXSIZE'] = '536870912' #512Mb (32mb default)
 
         os.environ['TF_MIN_GPU_MULTIPROCESSOR_COUNT'] = '2'
         os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' #tf log errors only