fix oom for batch_size=208 (#618)

2021-03-01 12:37:17 +08:00 · 2021-03-01 12:37:17 +08:00 · aeccae2128
parent ba17052a54
commit aeccae2128
1 changed files with 1 additions and 1 deletions
--- a/tools/static/train.py
+++ b/tools/static/train.py
@ -72,7 +72,7 @@ def main(args):
    if use_amp or use_pure_fp16:
        AMP_RELATED_FLAGS_SETTING = {
            'FLAGS_cudnn_exhaustive_search': 1,
-            'FLAGS_conv_workspace_size_limit': 4000,
+            'FLAGS_conv_workspace_size_limit': 1500,
            'FLAGS_cudnn_batchnorm_spatial_persistent': 1,
            'FLAGS_max_inplace_grad_add': 8,
        }