diff --git a/python/paddle/v2/fluid/layers/control_flow.py b/python/paddle/v2/fluid/layers/control_flow.py index 4b363ecbe78af82733fe1f80e44118a0dfda1f11..bef9602bb7cb94d45e663c7780a2b349c37a90b2 100644 --- a/python/paddle/v2/fluid/layers/control_flow.py +++ b/python/paddle/v2/fluid/layers/control_flow.py @@ -1220,7 +1220,8 @@ class DynamicRNN(object): self.lod_rank_table = None self.max_seq_len = None self.step_idx = None - self.zero_idx = fill_constant(shape=[1], value=0, dtype='int64') + self.zero_idx = fill_constant( + shape=[1], value=0, dtype='int64', force_cpu=True) self.mem_dict = dict() self.output_array = [] self.outputs = [] @@ -1275,7 +1276,8 @@ class DynamicRNN(object): def block(self): if self.status != DynamicRNN.BEFORE_RNN: raise ValueError("rnn.block() can only be invoke once") - self.step_idx = fill_constant(shape=[1], dtype='int64', value=0) + self.step_idx = fill_constant( + shape=[1], dtype='int64', value=0, force_cpu=True) self.step_idx.stop_gradient = False self.status = DynamicRNN.IN_RNN with self.while_op.block(): diff --git a/python/paddle/v2/fluid/layers/tensor.py b/python/paddle/v2/fluid/layers/tensor.py index 2608a8d1151fafa2da0cf7b605c4fa1210068057..2217c56b62af958709bb180bb360c938d1f0b312 100644 --- a/python/paddle/v2/fluid/layers/tensor.py +++ b/python/paddle/v2/fluid/layers/tensor.py @@ -180,7 +180,7 @@ def assign(input, output): return output -def fill_constant(shape, dtype, value, out=None): +def fill_constant(shape, dtype, value, force_cpu=False, out=None): """ **fill_constant** @@ -211,9 +211,12 @@ def fill_constant(shape, dtype, value, out=None): type='fill_constant', inputs={}, outputs={'Out': [out]}, - attrs={'shape': shape, - 'dtype': out.dtype, - 'value': float(value)}) + attrs={ + 'shape': shape, + 'dtype': out.dtype, + 'value': float(value), + 'force_cpu': force_cpu + }) out.stop_gradient = True return out