Update run_pretrain.py

835271fb · liangjing · fdac8fd9 · 835271fb
Commit 835271fb authored May 26, 2023 by liangjing
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

run_pretrain.py run_pretrain.py +2 -2

No files found.
--- a/run_pretrain.py
+++ b/run_pretrain.py
@@ -419,7 +419,7 @@ def dist_optimizer(args, optimizer):
    dist_strategy.fuse_grad_size_in_MB = 0
    if args.use_amp:
        dist_strategy.amp = True
-        custom_white_list = [
+        custom_black_list = [
          # 'softmax',
          # 'softmax_grad',
          # 'softmax_with_cross_entropy',
@@ -428,7 +428,7 @@ def dist_optimizer(args, optimizer):
          #  'reduce_sum_grad',
        ]

-        custom_black_list = [
+        custom_white_list = [
            'arg_max',
            'concat',
            'cumsum',