Fix the other lr_decay

Author: JiayiFeng

python/paddle/fluid/layers/learning_rate_scheduler.py

@@ -62,10 +62,10 @@ def noam_decay(d_model, warmup_steps):
         The decayed learning rate.
     """
     global_step = _decay_step_counter(1)
-    with init_on_cpu():
-        a = global_step**-0.5
-        b = (warmup_steps**-1.5) * global_step
-        lr_value = (d_model**-0.5) * ops.elementwise_min(a, b)
+
+    a = global_step**-0.5
+    b = (warmup_steps**-1.5) * global_step
+    lr_value = (d_model**-0.5) * ops.elementwise_min(a, b)
 
     return lr_value
 
@@ -108,12 +108,10 @@ def exponential_decay(learning_rate, decay_steps, decay_rate, staircase=False):
     """
     global_step = _decay_step_counter()
 
-    with init_on_cpu():
-        # update learning_rate
-        div_res = global_step / decay_steps
-        if staircase:
-            div_res = ops.floor(div_res)
-        decayed_lr = learning_rate * (decay_rate**div_res)
+    div_res = global_step / decay_steps
+    if staircase:
+        div_res = ops.floor(div_res)
+    decayed_lr = learning_rate * (decay_rate**div_res)
 
     return decayed_lr
 
@@ -138,11 +136,10 @@ def natural_exp_decay(learning_rate, decay_steps, decay_rate, staircase=False):
     """
     global_step = _decay_step_counter()
 
-    with init_on_cpu():
-        div_res = global_step / decay_steps
-        if staircase:
-            div_res = ops.floor(div_res)
-        decayed_lr = learning_rate * ops.exp(-1 * decay_rate * div_res)
+    div_res = global_step / decay_steps
+    if staircase:
+        div_res = ops.floor(div_res)
+    decayed_lr = learning_rate * ops.exp(-1 * decay_rate * div_res)
 
     return decayed_lr
 
@@ -184,12 +181,11 @@ def inverse_time_decay(learning_rate, decay_steps, decay_rate, staircase=False):
     """
     global_step = _decay_step_counter()
 
-    with init_on_cpu():
-        div_res = global_step / decay_steps
-        if staircase:
-            div_res = ops.floor(div_res)
+    div_res = global_step / decay_steps
+    if staircase:
+        div_res = ops.floor(div_res)
 
-        decayed_lr = learning_rate / (1 + decay_rate * div_res)
+    decayed_lr = learning_rate / (1 + decay_rate * div_res)
 
     return decayed_lr
 
@@ -224,25 +220,22 @@ def polynomial_decay(learning_rate,
     """
     global_step = _decay_step_counter()
 
-    with init_on_cpu():
-        if cycle:
-            div_res = ops.ceil(global_step / decay_steps)
-            zero_var = tensor.fill_constant(
-                shape=[1], dtype='float32', value=0.0)
-            one_var = tensor.fill_constant(
-                shape=[1], dtype='float32', value=1.0)
-
-            with control_flow.Switch() as switch:
-                with switch.case(global_step == zero_var):
-                    tensor.assign(input=one_var, output=div_res)
-            decay_steps = decay_steps * div_res
-        else:
-            decay_steps_var = tensor.fill_constant(
-                shape=[1], dtype='float32', value=float(decay_steps))
-            global_step = ops.elementwise_min(x=global_step, y=decay_steps_var)
-
-        decayed_lr = (learning_rate - end_learning_rate) * \
-                     ((1 - global_step / decay_steps) ** power) + end_learning_rate
+    if cycle:
+        div_res = ops.ceil(global_step / decay_steps)
+        zero_var = tensor.fill_constant(shape=[1], dtype='float32', value=0.0)
+        one_var = tensor.fill_constant(shape=[1], dtype='float32', value=1.0)
+
+        with control_flow.Switch() as switch:
+            with switch.case(global_step == zero_var):
+                tensor.assign(input=one_var, output=div_res)
+        decay_steps = decay_steps * div_res
+    else:
+        decay_steps_var = tensor.fill_constant(
+            shape=[1], dtype='float32', value=float(decay_steps))
+        global_step = ops.elementwise_min(x=global_step, y=decay_steps_var)
+
+    decayed_lr = (learning_rate - end_learning_rate) * \
+        ((1 - global_step / decay_steps) ** power) + end_learning_rate
     return decayed_lr