Commit 835271fb authored by liangjing's avatar liangjing
Browse files

Update run_pretrain.py

parent fdac8fd9
...@@ -419,7 +419,7 @@ def dist_optimizer(args, optimizer): ...@@ -419,7 +419,7 @@ def dist_optimizer(args, optimizer):
dist_strategy.fuse_grad_size_in_MB = 0 dist_strategy.fuse_grad_size_in_MB = 0
if args.use_amp: if args.use_amp:
dist_strategy.amp = True dist_strategy.amp = True
custom_white_list = [ custom_black_list = [
# 'softmax', # 'softmax',
# 'softmax_grad', # 'softmax_grad',
# 'softmax_with_cross_entropy', # 'softmax_with_cross_entropy',
...@@ -428,7 +428,7 @@ def dist_optimizer(args, optimizer): ...@@ -428,7 +428,7 @@ def dist_optimizer(args, optimizer):
# 'reduce_sum_grad', # 'reduce_sum_grad',
] ]
custom_black_list = [ custom_white_list = [
'arg_max', 'arg_max',
'concat', 'concat',
'cumsum', 'cumsum',
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment