From 318dd24cf3c7788f3fe88c7ccb910a09e6f469e5 Mon Sep 17 00:00:00 2001 From: Roc Date: Mon, 10 May 2021 11:10:03 +0800 Subject: [PATCH] sync_batch_norm supports amp --- paddleseg/core/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/paddleseg/core/train.py b/paddleseg/core/train.py index 151a95adeb..a582c8d2bb 100644 --- a/paddleseg/core/train.py +++ b/paddleseg/core/train.py @@ -157,7 +157,7 @@ def train(model, if fp16: with paddle.amp.auto_cast( - enable=True, custom_white_list={"elementwise_add", "batch_norm"}, custom_black_list={'bilinear_interp_v2'}): + enable=True, custom_white_list={"elementwise_add", "batch_norm", "sync_batch_norm"}, custom_black_list={'bilinear_interp_v2'}): if nranks > 1: logits_list = ddp_model(images) else: