diff --git a/python/paddle/v2/fluid/registry.py b/python/paddle/v2/fluid/registry.py index 6c0c3a35185391873fe5bb98d1ed5ee1cf13aa15..ff10542d40aabaf31897842754d38b7868472b21 100644 --- a/python/paddle/v2/fluid/registry.py +++ b/python/paddle/v2/fluid/registry.py @@ -1,16 +1,16 @@ # Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserve. # -#Licensed under the Apache License, Version 2.0 (the "License"); -#you may not use this file except in compliance with the License. -#You may obtain a copy of the License at +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # -#Unless required by applicable law or agreed to in writing, software -#distributed under the License is distributed on an "AS IS" BASIS, -#WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -#See the License for the specific language governing permissions and -#limitations under the License. +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. import re import cStringIO import warnings @@ -167,13 +167,18 @@ def register_layer(op_type): inputs[ipt.name] = val outputs = dict() - out = helper.create_tmp_variable(dtype=dtype) - outputs[o_name] = [out] + out = kwargs.pop(_convert_(o_name), []) + if out: + out_var = out[0] if (isinstance(out, list) or + isinstance(out, tuple)) else out + else: + out_var = helper.create_tmp_variable(dtype=dtype) + outputs[o_name] = [out_var] for name in intermediate_output_names: outputs[name] = [helper.create_tmp_variable(dtype=dtype)] helper.append_op( type=op_type, inputs=inputs, outputs=outputs, attrs=kwargs) - return helper.append_activation(out) + return helper.append_activation(out_var) func.__name__ = op_type func.__doc__ = _generate_doc_string_(op_proto)