From cfae890ca003c44b602b32d6c2f58b4c92fefcf4 Mon Sep 17 00:00:00 2001 From: yoonlee666 Date: Fri, 19 Jun 2020 20:28:01 +0800 Subject: [PATCH] support bert large pretrain with 96 bs --- model_zoo/bert/run_pretrain.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model_zoo/bert/run_pretrain.py b/model_zoo/bert/run_pretrain.py index 6abe9e9e83..65768946c1 100644 --- a/model_zoo/bert/run_pretrain.py +++ b/model_zoo/bert/run_pretrain.py @@ -81,7 +81,7 @@ def run_pretrain(): args_opt = parser.parse_args() context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target, device_id=args_opt.device_id) context.set_context(reserve_class_name_in_scope=False) - + context.set_context(variable_memory_max_size="30GB") ckpt_save_dir = args_opt.save_checkpoint_path if args_opt.distribute == "true": if args_opt.device_target == 'Ascend':