File tree Expand file tree Collapse file tree 2 files changed +10
-1
lines changed Expand file tree Collapse file tree 2 files changed +10
-1
lines changed Original file line number Diff line number Diff line change @@ -411,7 +411,7 @@ def _create(
411
411
config .enable_mkldnn ()
412
412
if "bf16" in self ._option .run_mode :
413
413
config .enable_mkldnn_bfloat16 ()
414
- config .set_mkldnn_cache_capacity (- 1 )
414
+ config .set_mkldnn_cache_capacity (self . _option . mkldnn_cache_capacity )
415
415
else :
416
416
logging .warning (
417
417
"MKL-DNN is not available. We will disable MKL-DNN."
Original file line number Diff line number Diff line change @@ -119,6 +119,7 @@ def _get_default_config(self):
119
119
"trt_dynamic_shape_input_data" : None , # only for trt
120
120
"trt_shape_range_info_path" : None , # only for trt
121
121
"trt_allow_rebuild_at_runtime" : True , # only for trt
122
+ "mkldnn_cache_capacity" : 10 ,
122
123
}
123
124
return default_config
124
125
@@ -294,6 +295,14 @@ def trt_allow_rebuild_at_runtime(self):
294
295
def trt_allow_rebuild_at_runtime (self , trt_allow_rebuild_at_runtime ):
295
296
self ._update ("trt_allow_rebuild_at_runtime" , trt_allow_rebuild_at_runtime )
296
297
298
+ @property
299
+ def mkldnn_cache_capacity (self ):
300
+ return self ._cfg ["mkldnn_cache_capacity" ]
301
+
302
+ @mkldnn_cache_capacity .setter
303
+ def mkldnn_cache_capacity (self , capacity : int ):
304
+ self ._update ("mkldnn_cache_capacity" , capacity )
305
+
297
306
# For backward compatibility
298
307
# TODO: Issue deprecation warnings
299
308
@property
You can’t perform that action at this time.
0 commit comments