diff --git a/tests/ut/python/ops/test_momentum.py b/tests/ut/python/ops/test_momentum.py
index 28b963701508c5af7abc19ad94a3a0d91c91d5d3..f25e4faf2d15542a0d08719c82a6c4052a9a0364 100644
--- a/tests/ut/python/ops/test_momentum.py
+++ b/tests/ut/python/ops/test_momentum.py
@@ -38,7 +38,7 @@ def tensor_run_opt(opt, iters, learning_rate, momentum,
                    gradient, variable, moment):
     """ tensor_run_opt """
     success = True
-    new_weight = opt(variable, moment, learning_rate, gradient, momentum)
+    new_weight = opt(variable, moment, learning_rate, gradient, momentum)[0]
     success = F.depend(success, F.assign(variable, new_weight))
     return success