Update benchmark_utils.py

zheyuye · Aug 10, 2020 · f0ae933 · f0ae933
1 parent 9735edb
commit f0ae933
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/scripts/benchmarks/benchmark_utils.py b/scripts/benchmarks/benchmark_utils.py
@@ -778,8 +778,8 @@ def train_step():
                     # contextual_embedding to contextual_embedding_ograd
                     # and the head gradient of pooled_out to pooled_out_ograd
                     # Thus, we simply doing two hadamard product and sum up the results.
-                    fake_loss = (contextual_embedding * contextual_embedding_ograd).sum()\
-                                + (pooled_out * pooled_out_ograd).sum()
+                    fake_loss = mxnet.np.sum(contextual_embedding * contextual_embedding_ograd)\
+                                + mxnet.np.sum(pooled_out * pooled_out_ograd)
                     fake_loss.backward()
                 mxnet.npx.waitall()
         elif model_cls.__name__ in ['BartModel']: