fix

Oneflow-Inc · Oct 7, 2020 · 0f617c2 · 0f617c2
1 parent 539750e
commit 0f617c2
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/LanguageModeling/BERT/pretrain.py b/LanguageModeling/BERT/pretrain.py
@@ -82,9 +82,9 @@ def PreTrain(
         initializer_range=initializer_range,
     )
     with flow.scope.namespace("cls-loss"):
-        if use_fp16:
-            lm_loss = flow.reduce_mean(lm_loss)
-            ns_loss = flow.reduce_mean(ns_loss)
+        if not use_fp16:
+            lm_loss = flow.math.reduce_mean(lm_loss)
+            ns_loss = flow.math.reduce_mean(ns_loss)
         total_loss = lm_loss + ns_loss
     return total_loss, lm_loss, ns_loss